怎样优化混合精度训练法

混合精度训练（Mixed Precision Training）通过结合 FP16/BF16 与 FP32 来加速训练、节省显存，但如果使用不当，也会导致精度下降、溢出（NaN/Inf）、收敛变慢等问题。下面从原理 → 常见问题 → 优化策略系统说明如何优化。

一、混合精度训练基本原理

• FP16/BF16：计算快、显存小（主要用于前向/反向）
• FP32：数值稳定（主要用于：
• 主权重（Master Weights）
• 优化器状态
• 易溢出的累加操作）

核心目标：

在几乎不损失精度的前提下，最大化速度与显存收益

二、常见性能与精度问题

1. 梯度下溢（Underflow）

• FP16 最小正值 ~6e-8，小梯度直接变 0

1. 梯度上溢（Overflow）

• 大激活 / 大学习率导致 Inf/NaN

1. Loss Scale 设置不合理

• 太大 → 溢出
• 太小 → 无效

1. 某些层不适合低精度

• LayerNorm、Softmax、Attention 数值敏感

1. 硬件未充分利用

• Tensor Core 未触发

三、核心优化策略（重点）

1. 使用成熟的混合精度框架

不要手写，优先使用：

• PyTorch
• torch.cuda.amp
• torch.nn.Transformer + fp16
• AMP（Automatic Mixed Precision）
• DeepSpeed / Megatron-LM
• NVIDIA Apex（旧项目）

✅ 好处：自动 cast、loss scaling、op 选择

2. 合理设置 Loss Scaling（FP16 关键）

动态 Loss Scale（推荐）

scaler = torch.cuda.amp.GradScaler()

• 自动翻倍 / 减半
• 溢出时跳过 update

手动调参建议

• 初始：2^15 ~ 2^16
• 观察：
• 频繁 skip → 降 scale
• 梯度长期很小 → 升 scale

⚠️ BF16 通常不需要 Loss Scale

3. 优先使用 BF16（如果硬件支持）

✅ 支持设备：

• Ampere 及以上（A100 / H100 / 3090 / 4090）

优势：

• 动态范围 ≈ FP32
• 几乎不溢出
• 不需要 loss scale

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
...

4. 对数值敏感模块强制 FP32

典型模块：

• LayerNorm
• Softmax
• LogSoftmax
• 小 head 的 classifier
• 累加型操作（sum/mean）

x = x.float()
x = F.layer_norm(x, ...)
x = x.half()

5. 优化器与状态精度管理

推荐方案

• Master Weight：FP32
• Optimizer State：FP32
• 梯度：FP16/BF16

DeepSpeed 示例：

{
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000
}
}

或 BF16：

"bf16": { "enabled": true }

6. 激活值显存优化（配合混合精度）

混合精度 + 以下技术效果最佳：

• Gradient Checkpointing
• Flash Attention
• Tensor Parallel / ZeRO

7. 避免常见反模式

❌ 所有层统一 FP16
❌ Loss Scale 写死
❌ 在 CPU 和 GPU 间频繁 cast

❌ 小 batch + 大模型 + FP16

四、工程级优化建议（大模型训练）

五、调试 checklist

• [ ] 是否出现 NaN/Inf
• [ ] 是否频繁 skip step
• [ ] Loss 是否震荡
• [ ] 是否启用 Tensor Core
• [ ] 是否 BF16 可用

六、总结一句话

优先 BF16 + AMP，敏感层 FP32，大模型配 ZeRO，FP16 必配 Loss Scale
如果你愿意，我可以：

• 给你 PyTorch 完整混合精度模板
• 针对 LLM / CNN / 多卡训练 给具体配置
• 帮你 排查 NaN / 精度下降问题

直接说你的模型和应用场景即可。