混合精度训练如何实现

混合精度训练（Mixed Precision Training）是指在神经网络训练过程中，同时使用单精度（FP32）和半精度（FP16 / BF16），以加快计算速度、减少显存占用，同时通过一些技巧保持模型收敛性和稳定性。

下面从 原理 → 关键问题 → 实现方式 → 框架示例 四个层面说明。

一、为什么要用混合精度

1. 优势

• 计算更快：GPU（如 NVIDIA Tensor Core）对 FP16 计算有数倍加速
• 显存更省：FP16 只占 FP32 一半空间
• 带宽更低：参数和激活传输更快

2. 问题

• 数值范围小：FP16 易溢出（INF / NaN）
• 舍入误差大：小梯度可能“下溢”为 0
• 累积误差：多次加法用 FP16 会不准

二、混合精度的核心思想

用 FP16 做“计算”，用 FP32 做“存储和累加”
典型策略：

• 权重副本：FP32 主权重（master weight）
• 前向 / 反向：FP16
• 梯度 / 更新：FP32

三、关键技术（必须做）

1. 损失缩放（Loss Scaling）【FP16 必须】

问题：小梯度在 FP16 中变为 0
解决：

loss * scale → backward → gradient / scale

• scale 初期较大（如 2^15）
• 若出现 NaN，自动减小 scale

2. Master Weights（主权重）

• 优化器更新使用 FP32 权重
• 每步将 FP32 权重转 FP16 用于前向

3. 算子精度选择

• ✅ 用 FP16：卷积、MatMul、激活
• ❌ 用 FP32：
• Softmax
• LayerNorm
• Loss
• 小型 reduce

4. BF16 的替代方案（更省心）

• BF16：范围同 FP32，精度较低
• 不需要 loss scaling
• 更适合 Transformer

四、实现方式（以框架为例）

1. PyTorch（推荐 AMP）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
optimizer.zero_grad()
with autocast():
output = model(input)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 自动：

• 插入 cast
• loss scaling
• 防止 NaN

2. PyTorch 手动混合精度（简化版）

model.half()
input = input.half()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
with torch.no_grad():
for p in model.parameters():
p.data = p.data - lr * p.grad.data

⚠ 不推荐（不稳定）

3. TensorFlow / Keras

tf.keras.mixed_precision.set_global_policy('mixed_float16')

• 自动 loss scaling
• 支持 GPU / TPU

4. DeepSpeed / Megatron

• FP16 / BF16
• ZeRO + 混合精度
• 大模型训练标配

五、典型训练流程（总结）

FP32 权重
↓
转 FP16
↓
前向（FP16）
↓
Loss × scale
↓
反向（FP16）
↓
梯度 ÷ scale
↓
FP32 更新权重

六、什么时候不适合

• 数值极敏感任务（如物理仿真）
• 极小模型（加速不明显）
• 旧 GPU（无 Tensor Core）

如果你愿意，我可以：

• 给你画一张 混合精度数据流图
• 对比 FP16 vs BF16
• 讲 大模型（LLM）中的混合精度实践

你更想看哪一个？