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

混合精度训练（Mixed Precision Training，通常用 FP16/BF16 + FP32）能显著提升训练速度和显存利用率，但也引入了一些数值稳定性和工程实现上的坑。下面从原理 → 注意事项 → 实践建议三个层面说明。

一、核心原理简述

FP16：计算快、显存小，但动态范围小（易溢出 / 下溢）
BF16：动态范围同 FP32，但精度低（更安全但效果依赖硬件）
FP32：用于“主权重”和“易出错计算”，保证稳定性

常见策略：前向/反向用低精度，更新用高精度（Master Weights）

二、混合精度训练的注意事项

1️⃣ 数值溢出（Overflow / Underflow）

FP16 最大值为 65504，稍大梯度就会 NaN
小梯度可能变成 0（下溢）

✅ 解决：

使用 Loss Scaling（尤其 FP16）
或直接使用 BF16（推荐在 A100/H100 上）

2️⃣ Loss Scaling（FP16 必备）

将 loss 乘以一个 scale 因子，使梯度保持在 FP16 可表示范围
反向后再除回来

⚠️ 注意：

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

✅ 现代框架（AMP）已自动管理：

torch.cuda.amp.GradScaler()

3️⃣ Master Weights（主权重）

优化器更新必须基于 FP32 权重
否则精度累积误差会被放大

✅ 框架通常自动维护：

FP16 权重（计算）
FP32 权重（更新）

4️⃣ 不适合低精度的操作

以下操作 不建议用 FP16/BF16：

softmax（尤其长序列）
layer norm / batch norm
exp / log / pow
小数值累加（sum reduction）

✅ 做法：

强制使用 FP32（autocast 已处理大部分）

5️⃣ 学习率与 Batch Size

混合精度常配合 更大 batch
学习率可能需要 线性缩放

⚠️ 不要直接照搬 FP32 的超参

6️⃣ 硬件与框架支持

硬件	建议
V100	FP16 + Loss Scaling
A100/H100	BF16（首选）
CPU	不推荐

框架：

PyTorch AMP ✅
TensorFlow MixedPrecisionPolicy ✅
手动实现 ❌（容易出错）

7️⃣ 收敛性与精度观察

混合精度可能：

收敛稍慢
最终精度略有差异（通常 < 0.1%）

✅ 建议：

前几个 epoch 对比 FP32
监控 loss / grad norm

三、PyTorch 示例（推荐写法）

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

四、一句话总结

混合精度训练不是“打开开关就行”，核心是：防止数值不稳定 + 保持主权重精度 + 合理使用 BF16/FP16
如果你愿意，我可以：

对比 FP16 vs BF16
针对 Transformer / LLM 给具体建议
帮你检查现有训练代码是否安全