混合精度训练注意什么 - GPU

混合精度训练（Mixed Precision Training，通常用 FP16/BF16 + FP32 结合）能显著加速训练、节省显存，但如果使用不当也会导致精度下降、梯度消失/溢出、收敛异常等问题。下面从实践角度总结注意事项。

一、核心原理回顾

FP32：精度高，用于存主权重、梯度、优化器状态
FP16/BF16：计算快、显存小，用于前向/反向计算
常见方案：
FP16 + Loss Scaling
BF16（推荐，几乎不用 loss scale）
AMP（PyTorch torch.cuda.amp）

二、关键注意事项

1. 使用 Loss Scaling（FP16 必做）

FP16 动态范围小，小梯度会下溢为 0
做法：
前向前放大 loss
反向后缩小梯度
PyTorch 示例：

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

✅ BF16 一般不用 loss scaling

2. 权重 & 优化器状态保持 FP32

不要将：
模型参数
优化器状态（Adam momentum / variance）

存成 FP16

否则：
更新量被截断
训练不收敛

✅ AMP 默认已处理

3. 警惕数值溢出（Overflow）

表现：

loss 变成 inf / nan
梯度爆炸

应对：

减小学习率
使用梯度裁剪
降低 loss scale 上限
换 BF16

4. 某些层不适合混合精度

谨慎使用 FP16 的层：

Softmax（尤其长序列）
LayerNorm / BatchNorm（统计量）
指数函数、log
小数值累加（sum reduction）

✅ 可手动指定这些层用 FP32

5. 注意 Batch Size 与显存

混合精度 ≠ 一定 2 倍 batch
实际加速受：
计算瓶颈
通信（多卡）
优化器类型

6. 分布式训练特别小心

DDP + AMP 要同步 GradScaler
梯度 all-reduce 前不能 scale 错
多卡建议：
BF16
或统一 loss scale 策略

7. 初始化与学习率

混合精度下：
学习率可稍大
但 warmup 更重要
推荐：
linear warmup
cosine decay

8. 评估 / 推理阶段

推理可用 FP16/BF16
但：
最终指标评估建议 FP32
尤其小数据集

9. 框架与硬件支持

硬件	建议
A100 / H100	BF16 最佳
V100	FP16 + AMP
T4	FP16 谨慎
CPU	基本无意义

三、常见坑总结

❌ loss 突然变 nan
❌ 训练初期不下降
❌ 精度明显低于 FP32
❌ 优化器状态被转成半精度
✅ 正确姿势：

主权重 FP32
计算半精度
梯度可控
监控 loss / grad norm

四、一句话建议

能用 BF16 就用 BF16；用 FP16 一定要 Loss Scaling，并保持优化器 FP32。
如果你用的是 PyTorch / TensorFlow / 大模型训练（LLM），我可以给你更具体的配置示例。