混合精度训练#

1. 先区分精度与范围#

全 FP32 训练把计算和状态都放在较高精度，成本较大。混合精度按算子和状态分别选择 dtype，在支持的硬件上减少计算或存储成本，同时为敏感计算保留精度。

FP16 尾数更长，但范围较小；BF16 范围接近 FP32，舍入更粗。TF32 是部分硬件执行 FP32 矩阵运算时的计算模式，不能把它当作与上述三者同类的 19 位张量存储类型。更多格式见低比特数值格式。

2. Loss scaling 解决哪一类问题#

FP16 最小正规正数约为 $6.10\times10^{-5}$，最小次正规正数约为 $5.96\times10^{-8}$。两者不同；次正规数还可能受到硬件冲零行为影响。梯度过小可能下溢，过大则可能上溢。Loss scaling 主要尝试保留反向中的小梯度，不能修复前向中已经出现的 Inf/NaN。

2.1 放大后为什么可以还原#

设损失为 $\mathcal L$，缩放因子 $S>0$ 在本次反向中固定。链式法则给出：

$$\nabla_\theta(S\mathcal L)=S\nabla_\theta\mathcal L,\qquad g=\frac{1}{S}\nabla_\theta(S\mathcal L).$$

例如真实梯度 $10^{-8}$ 可能在 FP16 中归零；取 $S=65536$，放大后的量为 $6.5536\times10^{-4}$，更容易表示。解除缩放与高精度更新仍有舍入误差，这不是严格恢复所有已损失信息。

2.2 为什么需要动态调整#

缩放太大可能使梯度溢出。动态 scaler 检查梯度：发生非有限值时跳过更新并减小 $S$；连续成功一定次数后再尝试增大。缩放在一个累积窗口内保持不变，具体调用只在有效更新边界发生，见 AMP examples。

3. 计算、累积和状态要分别核对#

Autocast 不会自动把整个模型和优化器状态都压成半精度。常规 AMP 可以保留 FP32 参数，部分算子使用低精度；不同框架的主副本与梯度存储也不同。因此“开启 AMP 后训练显存减半”没有通用保证。

BF16 通常无需 loss scaling，但仍有舍入、上溢和下溢边界。FP16 经常配合 scaler；选择应依据实际数值范围与硬件支持。不能由指数位相同推出 BF16 训练一定稳定。

4. 在有效更新边界使用 AMP#

一次更新的顺序为前向 autocast → 缩放损失并反向 → unscale → 裁剪 → scaler.step → scaler.update。完整代码与尾部累积处理集中在训练循环，本篇只维护数值机制。

代码采用 PyTorch 2.9 的 torch.autocast("cuda", ...) 与 torch.amp.GradScaler("cuda", ...) 接口；自定义算子的 dtype 需另行检查，不能假定所有 softmax 或归一化调用都自动得到同样处理。

参考资料#

• PyTorch 2.9 AMP：算子精度与 autocast/scaler 接口。
• Mixed Precision Training，§3：高精度权重副本、loss scaling 与算术精度分工。