如何调试混合精度训练

调试混合精度训练（Mixed Precision Training，常见如 PyTorch AMP / TensorFlow mixed precision）通常比纯 float32 训练更棘手，因为问题不一定直接报错，而表现为精度下降、loss 震荡、NaN/Inf、训练不收敛或结果异常。下面从「常见坑 → 调试方法 → 工具与实践」系统讲一下。

一、混合精度训练常见的问题来源

1. 数值下溢 / 上溢（最重要）

• float16 范围：约 6e-5 ~ 6e4
• 常见爆炸点：
• softmax / log_softmax 输入过大
• 除以很小的数（如小 batch 的 norm）
• 梯度非常小（下溢为 0）
• 指数运算（exp / pow）

2. 不安全的 op 被自动放进 float16

• 某些 op 在 float16 下不稳定：
• sum、mean 在大 tensor 上
• softmax（建议 float32）
• layernorm / batchnorm 统计量
• cross entropy（建议 float32）

3. Loss Scaling 问题（PyTorch AMP）

• scale 太大 → 梯度 Inf
• scale 太小 → 梯度为 0
• scale 动态下降太快 → 训练停滞

4. 模型结构或算法本身不适合 AMP

• 对数值极敏感的任务（如PINN、物理仿真）
• 极小学习率 + 极小梯度
• 自定义 CUDA op 未支持 half

二、系统化调试流程（强烈推荐）

✅ Step 1：先确认「float32 一定是对的」

这是底线

• 用 完全相同的 seed / batch / 数据
• 关闭 AMP，纯 float32 跑 1–2 epoch
• 保存：
• loss 曲线
• 某层权重 / 梯度统计

如果 float32 都不稳 → 先别用 AMP

✅ Step 2：开启 AMP，但“保守”

PyTorch 示例（推荐写法）

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

先观察：

• loss 是否 NaN
• scale 是否在下降

✅ Step 3：定位「哪里炸了」

方法 1：打印 NaN / Inf

for name, p in model.named_parameters():
if torch.isnan(p.grad).any():
print(name, "grad nan")

方法 2：逐层检查数值范围

def debug_tensor(x, name):
print(name, x.dtype, x.min(), x.max(), x.mean())

重点看：

• attention logits
• softmax 前
• norm 的 variance

✅ Step 4：强制关键 op 用 float32（最有效）

PyTorch 手动指定

x = x.float()
x = torch.softmax(x, dim=-1)
x = x.half()

或：

with torch.cuda.amp.autocast(enabled=False):
loss = criterion(logits, labels)

常见“必须 float32”的地方

• softmax / log_softmax
• cross entropy
• layer norm 统计
• 小数值除法

✅ Step 5：调整 Loss Scaling（PyTorch）

GradScaler(
init_scale=1024,
growth_factor=2,
backoff_factor=0.5,
growth_interval=2000
)

调试策略：

• NaN 多 → 减小 init_scale
• 训练慢 → 增大 growth_interval

三、TensorFlow / JAX 的注意点

TensorFlow

opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)

常见问题：

• custom training loop 忘记 get_scaled_loss
• dtype policy 不一致

JAX

• 显式 jax.numpy.float16
• 对 softmax 使用 float32 再转回

四、高级调试技巧

1. 对比 float32 vs float16 的梯度

grad_fp32 = grad(float32_model)
grad_fp16 = grad(amp_model)

看相对误差是否 > 10%

2. 小 batch 复现

• 用 2–4 条样本
• 禁用 shuffle
• 对比每一步数值

3. 使用 torch.autograd.detect_anomaly()

with torch.autograd.detect_anomaly():
loss.backward()

五、什么时候「不要用」混合精度

❌ 以下情况慎用：

• 数值极敏感（微分方程、量化前训练）
• 极小模型（收益 < 成本）
• 自研 op 不支持 half
• 调试阶段（先 float32）

六、一句话总结

混合精度调试的核心不是“开 AMP”，而是：
确保数值安全 op 用 float32 + 动态观察 loss scaling + 逐层定位 NaN 来源
如果你愿意，可以把：

• 框架（PyTorch / TF / JAX）
• 模型类型（Transformer / CNN / 自定义）
• 报错或现象（NaN / 不收敛）

发给我，我可以给你针对性调试方案。