踩坑实录：大模型混合精度（AMP）训练中的数值下溢与不稳定排查指南

在追求集群吞吐与显存利用率的过程中，混合精度训练（AMP）是标配方案。从理论上看，FP16 将计算吞吐翻倍、显存砍半；然而实际工程落地中，它也是训练任务在半夜无预警崩溃、Loss 突变为 NaN 或梯度静默下溢的元凶。

排查此类问题最忌盲目猜测试错——一次大模型 Checkpoint 恢复与重新跑批的成本极高。结合生产环境中的排查经验，以下是一套结构化的数值不稳定性诊断与修复路径。

一、 常见异常表象与核心机制

混合精度训练的不稳定，本质上是 动态范围与精度折中 带来的副产物。FP16 的指数位只有 5 位，最小的正规数约为 $6.10 \times 10^{-5}$，最大值仅为 $65504$。

通常遇到的并非单一的程序报错，而是以下三类特征表现：

1. Loss Scale 持续探底：开启 Dynamic Loss Scaling 后，日志中显示 scale 因子频繁下调（例如从 $2^{16}$ 连续除以 2 直到跌为 1 或更小），步长无法正常更新，训练陷入停滞。
2. 梯度突变引发的数值截断（Underflow）：深层网络或注意力层回传的微小梯度落入 FP16 下限之下，被硬件强制刷为 0（Flush-to-Zero），导致底层参数无法获得有效更新。
3. 突发性 NaN 炸模：常见于训练推进到数千步之后，Grad Norm 毫无征兆地由 $1.0$ 附近飙升至几十甚至上百，紧接着下一步直接崩溃。

二、 定位路径：避免全局断点的轻量级监控

遇到不稳定时，直接在多机多卡训练脚本里无脑挂上 torch.autograd.set_detect_anomaly(True) 往往会导致训练极其缓慢甚至显存直接 OOM。更可行的诊断手段是在关键节点注入轻量级监控。

1. 抓取最先出现异常的算子层

不要等到 optimizer.step() 抛错再去查。可以利用 PyTorch 的 Hook 机制针对可疑模块（通常是 Attention 层与 Norm 层）挂载输出与梯度的极值检测：

def check_nan_hook(module, grad_input, grad_output):
    for i, g in enumerate(grad_input):
        if g is not None:
            if torch.isnan(g).any() or torch.isinf(g).any():
                print(f"[Alert] Layer: {module.__class__.__name__} grad_input[{i}] NaN/Inf detected!")
                # 记录此时的最大绝对值与均值分布
                print(f"Max: {g.abs().max().item()}, Non-zero ratio: {(g != 0).float().mean().item()}")

2. 区分是“下溢”还是“溢出导致的除以零”

观察 Tensor 的非零元素比例。如果某一层的梯度非零率（Sparsity）突然大幅攀升至接近 100% 且全为 0，基本判定为数值下溢被截断；若是出现 Inf 后引发连环 NaN，则是典型的激活值或乘积爆炸。

三、 高频踩坑点与根因分析

1. Softmax 与 Attention 矩阵计算

标准 Attention 中 $Q K^T / \sqrt{d}$ 的计算，在序列长度较长时，未经 Scaled 的点积极易超出 FP16 的最大上限 65504，导致后续 Softmax 直接输入 Inf。此外，如果 Attention Mask 采用类似 -1e9 的大负数来屏蔽 Padding，在 FP16 下该数值早已溢出（FP16 最小表示约为 -65504），结果直接变成 -inf 或未定义行为。

2. LayerNorm / RMSNorm 的方差统计

Normalization 层的方差计算涉及对输入特征的平方求和（$\sum x^2$）。当 Hidden Size 达到 4096 甚至更高时，累加很容易瞬间顶破 65504。

3. 优化器更新步长过小

Adam 优化器的一阶矩 $m_t$ 和二阶矩 $v_t$ 如果也保存在低精度中，参数实际更新量 $\Delta \theta = -\eta \frac{m_t}{\sqrt{v_t} + \epsilon}$ 会极度微小。一旦更新量小于 FP16 能表示的最小差值，参数在实际相加时就会发生数值吸收，权重彻底停滞不动。

四、 生产环境治理方案

针对上述根因，工程实践中建议按以下优先级落地改动：

1. 关键算子强制回退到 FP32（White-listing）

   • Norm 层与 Softmax：不要在 FP16 模式下直接计算方差与 Softmax。利用 PyTorch torch.cuda.amp.autocast(enabled=False) 或在算子内部显式转为 x.float()，计算完成后再转换回低精度。
   • Attention Mask 处理：Padding Mask 的占位值从 -1e9 改为 -1e4，保证在 FP16 的动态表示范围之内。

2. 转向 BF16（Bfloat16）体系 如果硬件支持（NVIDIA Ampere 架构如 A100/H100/RTX 30 系列及以上），优先放弃 FP16 全面迁移至 BF16。

   • BF16 拥有与 FP32 完全相同的 8 位指数位，虽然尾数精度略逊于 FP16，但彻底杜绝了动态范围不足引发的下溢与溢出问题。
   • 切至 BF16 后通常可以直接拿掉复杂且脆弱的 Dynamic Loss Scaler，训练稳定性会得到质的提升。

3. Loss Scale 与梯度裁剪（Gradient Clipping）协同 若必须受限使用 FP16：

   • 调大 Loss Scaler 的平稳窗口：将 growth_interval 从默认的 2000 步调高至更大步数，避免由于单步偶发噪点引起 Scale 震荡。
   • 裁剪阈值提前：在缩放还原（scaler.unscale_(optimizer)）后，立刻执行 torch.nn.utils.clip_grad_norm_，将全局梯度范数严格限制在 1.0 或更低，阻断梯度尖峰蔓延至权重更新。

排查混合精度的数值隐患，核心是守住动态范围的上下界。通过算子层面的精度回退兜底，辅以结构化的极值监控，才能在大模型长时间的连续训练中保障稳定性与算力产出比。