3、混合精度训练原理：自动混合精度（AMP）、损失缩放（Loss Scaling）、梯度裁剪与参数更新流程

混合精度训练并不是简单地把模型权重从 FP32 换成 FP16，而是围绕“哪些计算用低精度、哪些状态必须保留高精度、梯度如何安全回传”建立的一整套协同机制。理解 AMP、Loss Scaling、梯度裁剪与参数更新之间的先后关系，是后续排查精度异常和做参数校准的基础。

一、核心原理：AMP 的“双精度分工”

自动混合精度（AMP）的核心思想是：前向与反向的矩阵乘、卷积等计算密集型算子使用 FP16/BF16 加速，而参数主副本、优化器状态、归约与敏感算子保留 FP32。这样既获得 Tensor Core 的吞吐收益，又避免小梯度在低精度下下溢为零。

以 PyTorch AMP 为例，典型结构如下：

scaler = torch.cuda.amp.GradScaler()

for x, y in loader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(dtype=torch.float16):
        out = model(x)
        loss = criterion(out, y)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

这段代码里，autocast 负责算子级精度选择，GradScaler 负责损失缩放，unscale_ 负责把缩放后的梯度还原，clip_grad_norm_ 负责梯度裁剪，step 与 update 负责参数更新和缩放因子调整。四者顺序不能随意调换。

AMP 混合精度训练数据流与精度分工 输入 batch FP32 / FP16 autocast 前向 FP16 矩阵乘 / 卷积 loss 计算 FP32 累加更稳 scaler.scale(loss) 放大 2^S 倍 反向传播 FP16 梯度 unscale_ 梯度 ÷ 2^S clip_grad_norm_ max_norm 裁剪 scaler.step FP32 主权重更新 精度分工原则 • 计算密集算子：FP16/BF16，走 Tensor Core，吞吐高 • 参数主副本 / 优化器状态：FP32，避免累积误差 • 归约、Softmax、LayerNorm 等敏感算子：FP32 或 BF16 动态范围更安全

二、Loss Scaling：让梯度“浮”起来

FP16 的最小正规数约为 6.1e-5，很多层的梯度值远小于这个量级，直接反向会下溢为 0，导致参数不更新。Loss Scaling 的做法是：在反向传播前把 loss 乘以一个缩放因子 S，使梯度整体放大 2^S 倍，反向结束后再除以 S 还原。

动态损失缩放（Dynamic Loss Scaling）会根据梯度是否出现 inf/NaN 自动调整 S：

  • 连续 N 步没有溢出：S 翻倍，提升数值分辨率；
  • 出现 inf/NaN：跳过该步参数更新，S 减半，避免污染权重；
  • 缩放因子有上下界，防止无限增长或缩到 1 以下。
BF16 的指数位与 FP32 相同，动态范围足够大，通常不需要 Loss Scaling；FP16 则必须配合 Loss Scaling 才能稳定训练。

三、梯度裁剪与参数更新的正确顺序

梯度裁剪必须在 unscale_ 之后、scaler.step 之前执行。如果先裁剪再还原，裁剪阈值会被缩放因子放大，实际裁剪强度失真；如果先 step 再裁剪，则已经用未裁剪梯度更新了参数，裁剪失去意义。

步骤操作精度关键点
1autocast 前向FP16/BF16算子自动选择低精度
2loss 计算FP32避免累加误差
3scaler.scale(loss)FP32放大 2^S 倍
4backwardFP16梯度放大后回传
5unscale_FP32还原真实梯度
6clip_grad_norm_FP32按真实梯度裁剪
7scaler.stepFP32更新主权重
8scaler.update—调整 S

四、案例场景：一次典型的 AMP 精度异常

某视觉模型在开启 AMP 后，前 200 步 loss 正常下降，随后突然变为 NaN。排查日志发现：

[step 198] loss=0.421  scale=65536.0
[step 199] loss=0.418  scale=65536.0
[step 200] loss=nan    scale=32768.0  skipped=True
[step 201] loss=nan    scale=16384.0  skipped=True

根因是梯度裁剪写在 unscale_ 之前，导致裁剪阈值被放大 65536 倍，实际几乎没有裁剪，异常梯度直接进入优化器。调整顺序后，NaN 消失，训练恢复稳定。

常见误区：把 clip_grad_norm_ 放在 scaler.scale(loss).backward() 之后、unscale_ 之前；或忘记调用 unscale_ 就直接裁剪。这两种写法都会让裁剪失效。

五、复盘：参数校准的检查清单

  1. 确认 autocast 只包裹前向，不要包裹 loss 计算和优化器步骤；
  2. 确认 unscale_ 在裁剪之前调用，且每个 optimizer 只调用一次；
  3. 确认 clip_grad_norm_ 的 max_norm 与 FP32 训练时一致，不要因缩放而调大；
  4. 监控 scale 值变化，若持续下降说明梯度频繁溢出，需检查数据或模型数值稳定性；
  5. 若使用 BF16，可关闭 Loss Scaling，但仍需关注梯度裁剪与学习率匹配。
← 上一章 📖 返回目录 下一章 →