如何避免混合精度训练溢出

避免混合精度训练（通常主要是避免 FP16/半精度溢出）的关键在于：及时发现数值范围问题、控制梯度范围、并使用稳定的优化策略。下面从原理到实践系统说明。

一、混合精度训练为什么会溢出？

在混合精度训练中：

• 权重、激活、梯度：FP16（范围约 ±65504）
• 主权重（master weights）：FP32（用于参数更新）

常见溢出原因：

1. 梯度或激活值绝对值过大 → FP16 溢出 → 变为 inf / nan
2. Loss scaling 不当：

• 过小 → 下溢（变为 0）
• 过大 → 上溢

1. 优化器状态累积误差（尤其是 Adam）
2. 数值不稳定操作（如 softmax、log、除小值）

二、核心技术方案（必须掌握的）

✅ 1. 使用 Loss Scaling（最重要）

原理：
在反向传播前把 loss 放大，使梯度在 FP16 范围内；更新参数前再缩小。

# PyTorch 示例
from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()
for x, y in data:
with autocast():
y_pred = model(x)
loss = criterion(y_pred, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

✅ 注意事项：

• 不要手动设置 fixed scale
• 使用 动态 loss scaling（默认行为）

✅ 2. 保持 Master Weights（FP32）

• 前向 & 反向：FP16
• 参数更新：FP32 主权重

自动完成方式：

optimizer = torch.optim.Adam(model.parameters())
# + GradScaler（会自动维护 master weights）

⚠️ 不要自己手动把参数转成 FP16 再更新

✅ 3. 梯度裁剪（Gradient Clipping）

防止梯度爆炸导致溢出：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

搭配 AMP：

scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

✅ 推荐范围：

• NLP / Transformer：1.0
• CV / 小模型：5~10

三、模型与结构层面的防溢出技巧

✅ 4. 避免 softmax / log 的极端值

❌ 不推荐：

log_softmax(x)  # x 可能很大

✅ 推荐：

F.log_softmax(x, dim=-1)
# 或
F.cross_entropy(logits, target)

内部已实现数值稳定版本。

✅ 5. 谨慎处理除法与小数值

❌：

a / b  # b 很小

✅：

a / (b + eps)

或改用：

torch.clamp(b, min=eps)

✅ 6. LayerNorm > BatchNorm（在 FP16 下）

• LayerNorm：数值更稳定
• BatchNorm 在 FP16 下容易出现统计量溢出

Transformer 中几乎都使用 LayerNorm 也是这个原因。

四、训练策略层面的建议

✅ 7. 学习率要适当降低

混合精度 ≠ 可以随意增大学习率。
经验法则：

• 使用 AMP 后，学习率 保持或略小于 FP32

✅ 8. 在模型稳定后再启用 AMP

推荐流程：

1. 先用 FP32 跑几步，确保模型本身稳定
2. 再开启混合精度

✅ 9. 检测溢出并自动跳过更新

AMP 已内置：

scaler.update()

如果检测到 inf/nan：

• 自动跳过该 step
• 动态调整 loss scale

五、调试与监控技巧

✅ 10. 主动检测 NaN / Inf

for name, param in model.named_parameters():
if torch.isnan(param).any() or torch.isinf(param).any():
print(name)

或：

torch.autograd.detect_anomaly()

六、总结：防溢出 Checklist ✅

✅ 使用 torch.cuda.amp
✅ 动态 Loss Scaling
✅ 保持 FP32 master weights
✅ 梯度裁剪
✅ 避免极端数值操作
✅ 优先 LayerNorm
✅ 学习率不过大

✅ 先 FP32 验证模型稳定性

如果你愿意，我可以：

• ✅ 帮你 检查某段代码是否会溢出
• ✅ 针对 Transformer / CNN / 大模型 给出具体 AMP 配置
• ✅ 解释 为什么 BF16 比 FP16 更不容易溢出

只要告诉我你的模型类型即可。