混合精度训练如何调参

混合精度训练（Mixed Precision Training，通常指 FP16/BF16 + FP32）能显著加速训练并节省显存，但也会引入数值不稳定、收敛变慢或精度下降等问题。调参的核心目标是在保持收敛性和最终精度的前提下，尽可能享受加速收益。

下面按“从易到难”的顺序讲如何调参。

一、先确认基础设置是否正确

1. 选对混合精度类型

• FP16：常见于 NVIDIA GPU（V100/A100），需要小心溢出
• BF16：A100/H100/TPU 推荐，动态范围大，通常更稳
• 优先用框架自带方案：
• PyTorch：torch.cuda.amp（FP16）或 torch.autocast(dtype=bfloat16)
• TensorFlow：mixed_float16 policy
• DeepSpeed / Megatron：内置 AMP

✅ 经验：能 BF16 就 BF16，调参成本低很多。

二、最核心参数：Loss Scaling（FP16 才有）

2.1 什么是 Loss Scale

FP16 梯度太小会下溢为 0，所以：

scaled_loss = loss * loss_scale
grad = d(scaled_loss)
grad /= loss_scale

2.2 怎么调

• 动态 Loss Scale（推荐）
• 初始值：2^15 ~ 2^16
• 自动翻倍 / 减半
• 静态 Loss Scale
• 从小开始试：2^8 → 2^10 → 2^12
• 出现 NaN 就降，训练慢就升

✅ 实用建议：

• 一上来就 NaN → loss scale 太大
• loss 不下降 → loss scale 太小或 master weight 没用

三、Master Weights（必须开）

3. 保持 FP32 主权重

• 优化器状态（momentum / variance）必须是 FP32
• 参数更新在 FP32 中进行

PyTorch AMP 默认已做；
手动实现时要显式保留 model.float() 副本。

✅ 这是“不掉点”的关键，不要省。

四、学习率与优化器调参

4.1 学习率

混合精度不改变理论 LR，但：

• FP16 下梯度噪声更大
• 有时需要：
• LR × 0.8 ~ 1.0
• warmup 稍长一点（如 5% → 8%）

4.2 Optimizer

• Adam / AdamW：状态用 FP32
• 对 LR scheduler 无特殊要求
• 若出现震荡：
• 减小 LR
• 增大 eps（如 1e-4 → 1e-3）

五、数值不稳定时的“救命参数”

5.1 哪些层容易炸

• Softmax（尤其长序列）
• LayerNorm 前的 sum
• 大 logits（语言模型、CTC）
• 小初始化网络

✅ 对策：

• 对这些层强制 FP32：

with torch.autocast(enabled=False):
logits = model.head(x)

5.2 梯度裁剪

• FP16 下更容易梯度突变
• clip_grad_norm_(1.0) 通常够用
• 不稳定可试 0.5 ~ 2.0

六、Batch Size 与显存利用

混合精度省显存，通常可以：

• Batch size × 1.5 ~ 2
• 或加大序列长度

⚠️ 注意：

• Batch 变大 → LR 可能要线性 scaling
• 太大反而训练更噪

七、调试流程（推荐套路）

1. 小数据 + FP32 跑通
2. 切 BF16，不动超参
3. 再切 FP16：

• 开动态 loss scale
• 开 master weight

1. 看前 100 step：

• loss 曲线是否一致
• 有无 NaN / Inf

1. 不一致就：

• 降 loss scale
• 关键层回 FP32
• 加 warmup

八、常见坑速查表

如果你愿意，可以告诉我：

• 用的框架（PyTorch / TF / DeepSpeed）
• GPU 型号
• 任务类型（NLP / CV / 多模态）

我可以直接给你一套可抄的混精训练配置。