混合精度训练如何加速

混合精度训练（Mixed Precision Training）主要通过用低精度浮点数（如 FP16 / BF16）替代部分高精度计算（FP32）来加速模型训练，同时尽量保持数值稳定性和最终精度。下面从原理、加速点、关键技术和实践建议几个角度说明。

一、为什么能加速？

1. 计算更快（算力提升）

• GPU（如 NVIDIA Volta、Ampere、Hopper）对 FP16 / BF16 的吞吐通常是 FP32 的 2～8 倍
• 矩阵乘法（GEMM）、卷积等核心算子在低精度下并行度更高

2. 显存占用更少

• FP16 只占 2 字节，FP32 占 4 字节
• 可：
• 增大 batch size
• 使用更大模型
• 减少显存碎片和 OOM 风险

3. 带宽压力更小

• 激活值、梯度在显存和计算单元之间传输量减半
• 对大模型、长序列尤其明显

二、混合精度是怎么“混合”的？

典型策略（以 FP16 为例）：

BF16 不需要 loss scaling，动态范围与 FP32 一致，更稳定

三、关键技术（避免精度崩）

1. Loss Scaling（FP16 专用）

• 问题：小梯度在 FP16 下变成 0
• 解决：
• 前向前放大 loss
• 反向后缩小梯度

scaled_loss = loss * 1024
scaled_loss.backward()

2. Master Weights

• 优化器更新使用 FP32 权重副本
• 避免累积误差

3. 数值敏感操作保留 FP32

• softmax
• layer norm
• log / exp
• 小数值累加

四、框架层面的加速方式

PyTorch（AMP）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
out = model(x)
loss = criterion(out, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

TensorFlow

policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

五、典型加速效果

• 训练速度提升：1.5x ～ 3x
• 显存节省：30% ～ 50%
• 精度损失：通常 < 0.1%

六、不适合 / 需注意的场景

• 极度数值敏感任务（部分科学计算）
• 自定义 CUDA 算子不支持半精度
• 老旧 GPU（无 Tensor Core）

如果你愿意，我可以：

• 给你 某个模型（如 LLM / CNN）的混合精度配置
• 对比 FP16 vs BF16
• 讲 分布式 + 混合精度（FSDP / DeepSpeed)