4、动态损失缩放机制：损失缩放的目的、梯度下溢与上溢问题、动态调整策略
在混合精度训练中，FP16 的数值范围远小于 FP32。FP16 能表示的最小正规数约为 6.1×10⁻⁵，而 FP32 可低至约 1.2×10⁻³⁸。这意味着许多在 FP32 中正常存在的微小梯度，在 FP16 中会直接变成 0，导致参数无法更新。动态损失缩放正是为了解决这一矛盾而设计的核心机制。
一、损失缩放的目的
损失缩放的基本思想非常直接：在反向传播之前，把损失值乘以一个较大的缩放因子 S，使损失放大 S 倍。由于链式法则，所有梯度也会同比放大 S 倍，从而把原本会下溢到 0 的微小梯度“抬升”到 FP16 可表示的范围内。在参数更新前，再把梯度除以 S 还原回真实尺度。
这一过程可以用公式表达：
loss_scaled = loss * S
grad_scaled = backward(loss_scaled)
grad_real = grad_scaled / S
param -= lr * grad_real缩放因子 S 通常取 2 的幂次，如 2¹⁵、2¹⁶、2¹⁷，这样在乘除时只涉及指数位变化，不会引入额外舍入误差，计算开销也极低。
二、梯度下溢与上溢问题
损失缩放并非越大越好。缩放因子过小，梯度仍然下溢；缩放因子过大，梯度又会溢出为 Inf 或 NaN。这两种失败模式分别称为下溢和上溢。
- 下溢：梯度值小于 FP16 最小可表示范围，被截断为 0。表现为参数几乎不更新，loss 下降缓慢或停滞。
- 上溢：梯度值超过 FP16 最大值 65504，变成 Inf。一旦出现 Inf，反向传播中的 NaN 会迅速污染所有参数，训练直接崩溃。
因此，缩放因子必须落在一个“安全窗口”内：足够大以避免下溢，又足够小以避免上溢。而这个窗口会随着训练阶段、网络层、batch 组成动态变化，固定值很难长期有效。
关键判断：如果训练中出现 loss 突然变为 NaN，优先检查是否发生梯度上溢；如果 loss 长期不降且梯度范数接近 0，优先检查是否发生下溢。
三、动态调整策略
动态损失缩放的核心是：不预设固定 S，而是让 S 根据梯度是否溢出自动升降。主流框架（PyTorch AMP、NVIDIA Apex、TensorFlow mixed precision）都采用类似策略。
典型调整规则如下：
- 每隔 N 个 iteration 检查一次梯度中是否出现 Inf 或 NaN。
- 若连续若干步没有溢出，说明当前 S 偏保守，将 S 乘以 2（增大）。
- 若某一步检测到溢出，跳过该步参数更新，并将 S 乘以 0.5（减小）。
- 记录连续无溢出步数，达到阈值后再执行增大操作，避免频繁抖动。
PyTorch 的 torch.cuda.amp.GradScaler 封装了上述逻辑，使用方式如下：
scaler = torch.cuda.amp.GradScaler(
init_scale=2.**16,
growth_factor=2.0,
backoff_factor=0.5,
growth_interval=2000
)
for x, y in loader:
optimizer.zero_grad()
with torch.cuda.amp.autocast():
out = model(x)
loss = criterion(out, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()其中 init_scale 是初始缩放因子，growth_interval 是连续无溢出多少次后增大，backoff_factor 是溢出后的缩小比例。实际训练中，S 会在几千步内自动收敛到一个适合当前模型的稳定区间。
四、参数配置与场景案例
不同任务对缩放因子的敏感度不同。下表给出常见场景的参考配置：
| 场景 | 初始 S | growth_interval | 常见问题 |
|---|---|---|---|
| 小模型 / 小 batch | 2¹² ~ 2¹⁴ | 1000 | 下溢，loss 不降 |
| Transformer 大模型 | 2¹⁶ ~ 2¹⁸ | 2000 | 上溢，偶发 NaN |
| 目标检测 / 分割 | 2¹⁵ | 1500 | 梯度尺度差异大 |
| 微调预训练模型 | 2¹⁴ ~ 2¹⁶ | 2000 | 初期溢出频繁 |
一个典型排障案例：某 Transformer 训练在第 3000 步突然 NaN。日志显示 GradScaler 的 scale 从 65536 连续降到 8192，说明发生了多次上溢。将 init_scale 从 2¹⁶ 降到 2¹⁴，并把 growth_interval 从 1000 提高到 3000 后，训练稳定收敛。这说明初始 S 过大时，动态机制虽然能救回来，但频繁的 backoff 会浪费大量 step。
五、常见误区与复盘
- 误区一：认为动态缩放可以完全替代手动调参。实际上初始 S 和 growth_interval 仍需要根据任务设定，否则前期会浪费大量迭代。
- 误区二：看到 NaN 就只调小学习率。梯度上溢往往来自缩放因子过大，应先检查 scaler 的 scale 变化曲线。
- 误区三：忽略
scaler.step()与scaler.update()的配对。漏掉 update 会导致 S 永远不调整，退化为固定缩放。 - 误区四：在梯度累积场景中直接对每个 micro-batch 调用 update。正确做法是累积完成后统一 step 和 update，否则缩放因子会被错误地频繁调整。
调试建议：在训练循环中打印 scaler.get_scale()，观察 S 随 step 的变化曲线。稳定训练中 S 应在一个区间内小幅波动，而不是持续下降或持续上升。
总结来说，动态损失缩放是混合精度训练能否稳定收敛的关键开关。理解下溢与上溢的边界、掌握 scaler 的升降逻辑、并根据任务合理设置初始值和增长间隔，才能让 FP16 训练既快又稳。