一、混合精度训练的数学原理

说起混合精度训练，很多人觉得就是简单地把模型从 FP32 换成 FP16 跑。

理论上是这样，但真动手你会发现——数值范围一下就缩水了。

我最早接触 FP16 时，loss 直接炸成 NaN。排查半天，原来是一个梯度值超过了 65504。FP16 能表示的最大正数就是这么多。你想想看，FP32 能表示到 3.4e38，这一下差了 5 个数量级。

怎么办？加个缩放因子呗。

1.1 数值溢出问题与缩放因子的提出

FP16 的数值分布和 FP32 不一样。它只有 5 位指数位，10 位尾数位。这意味着它能表示的范围非常有限。

我做个小对比：

类型 指数位 尾数位 最大正数 最小正规格化数
FP32 8 23 3.4e38 1.18e-38
FP16 5 10 65504 6.10e-5

看到了吗？FP16 的最小正规格化数只有 6.10e-5。这就意味着，任何小于这个值的数都会被下溢成 0。上溢就更别提了，超过 65504 就变 Inf。

我在项目中遇到过，训练一个 LLM 时，激活值的分布特别宽。有 10% 的值都在 0.001 以下，直接就被 FP16 截断了。模型训练了三天，精度就是上不去。后来我们加了损失缩放，才把问题解决。

缩放因子的核心思路很简单：

  • 把计算图里的值整体放大（乘以一个缩放因子 S），让它落入 FP16 的安全范围
  • 等梯度传到优化器之前，再把缩放去掉（除以 S），保证权重更新正确

这个 S 就是我们要找的缩放因子。

1.2 梯度缩放原理与动态更新策略

梯度缩放是混合精度训练里最关键的数学技巧。说白了，它不是缩放模型参数，而是缩放损失值。

为什么缩放损失？因为梯度是通过反向传播从损失算出来的。你把损失放大 10 倍，反向传播算出来的梯度自然也会放大 10 倍。这样原本可能下溢的梯度，就被抬到了 FP16 能表示的范围里。

具体流程长这样：

  1. 前向传播：用 FP16 算模型输出
  2. 计算损失：损失函数输出 FP32 的标量 L
  3. 缩放损失：L_scaled = L × S（S 是缩放因子）
  4. 反向传播：在 L_scaled 上算梯度，所有中间梯度都是 FP16 格式
  5. 反缩放梯度：grad_fp32 = grad_fp16 / S，把梯度转回 FP32
  6. 权重更新：用 FP32 的优化器状态更新权重

举个实际例子。在 PyTorch 里，梯度缩放是这样做的：

scaler = torch.cuda.amp.GradScaler()

for data, label in dataloader:
    optimizer.zero_grad()
    
    with torch.cuda.amp.autocast():
        output = model(data)
        loss = loss_fn(output, label)
    
    # 缩放损失后进行反向传播
    scaler.scale(loss).backward()
    
    # 反缩放后更新权重
    scaler.step(optimizer)
    
    # 动态调整缩放因子
    scaler.update()

这里有个非常关键的点：缩放因子 S 不是人肉调的。模型在训练过程中，梯度分布会变化，你需要动态调整 S。

动态更新的逻辑：

  • 如果连续 N 步都没有出现 Inf/NaN（我一般设 N=2048），就把 S 增大 2x
  • 如果某一步出现了 Inf/NaN，说明 S 太大导致上溢了，就把 S 缩小 2x，然后跳过这步

我曾经在训练一个图像分割模型时，用了静态缩放因子 2^8，结果 loss 死活不收敛。后来改为动态缩放，S 在 2^10 到 2^14 之间来回跳，效果一下就好了。

这里有个警告，你们注意一下：

⚠️ 缩放因子和优化器不匹配
有些优化器（比如 LAMB）对梯度的大小特别敏感。你在缩放损失时放大梯度，可能导致优化器步长失准。我个人习惯是用 FP32 的 AdamW 配合梯度缩放，效果最稳定。如果你用了 LAMB 或 AdaBelief，一定要在验证集上仔细检查。

1.3 精度损失分析与实践考量

精度损失是混合精度训练绕不开的话题。很多人问我：到底损失了多少精度？

我一般从三个方面来拆解：

  • 上溢损失：梯度或激活值超过 FP16 上限，直接变为 Inf
  • 下溢损失：梯度值太小，在 FP16 下变成 0，优化器收不到更新信号
  • 舍入损失：FP16 尾数位只有 10 位，很多小数值被截断或四舍五入

先说上溢，这个最好防。动态缩放因子会自动调小 S，防止上溢。我训练 1000 个模型，大概只有 2-3 个会在训练初期出现持续性上溢，基本都是初始化有问题。

下溢就比较坑了。有些层的梯度天生就小，比如 embedding 层、BN 层的 gamma 梯度。我建议对这些层单独处理：

# 对特定层保持 FP32
model.embedding = model.embedding.float()
model.bn_layers = [layer.float() for layer in model.bn_layers]

# 或者手动调整缩放策略
scaler = GradScaler(init_scale=2**16, growth_interval=1000)

舍入损失是最隐形的。FP16 的最小精度是 2^-10 ≈ 0.001。也就是说，两个数相差如果小于 0.001，在 FP16 里就分不出来。这种情况在小梯度场景下特别明显。

我做过一个实验，对比不同缩放因子下的精度差异：

缩放因子 S 训练 Loss 验证集 AUC Inf/NaN 出现次数
2^8 0.423 0.934 12
2^12 0.418 0.937 3
动态缩放 0.415 0.939 1

你看，动态缩放效果最好。静态缩放 S=2^12 虽然 Inf 少，但舍入误差还是影响了精度。

💡 实践建议
1. 别手动设缩放因子，用自动动态缩放
2. 训练前 200 步密切关注 loss 曲线，如果突然飙升，先调小 init_scale
3. 把 embedding 层和 BN 层留在 FP32，损失很小但精度提升明显

精度损失这件事，其实没你想的那么可怕。我们现在训练 175B 的 LLM，全部用混合精度训练。只要缩放因子配得好，精度基本能打平纯 FP32。我甚至遇到过精度反超的情况——因为 FP16 训练更快，允许你在同样时间内遍历更多数据。

嗯，这就是混合精度训练的核心数学逻辑。缩放因子不是玄学，它是对抗 FP16 数值范围缺陷的数学补偿方案。你只要把握好上溢、下溢、舍入这三道坎，模型精度就不会垮。

核心总结
• 缩放因子解决数值溢出问题（上溢 → Inf，下溢 → 0）
• 损失缩放 + 动态更新策略，是混合精度训练的基石
• 精度损失主要来自 FP16 尾数位不足，通常可通过动态缩放弥补
• 特殊层（embedding、BN）建议保留 FP32 计算

下面这张图把这套流程串起来了：

混合精度训练梯度缩放流程 FP16 前向传播 计算 Loss (FP32) 缩放损失 L × S FP16 反向传播 反缩放梯度 ÷ S 动态调整 S 有 Inf 缩小，无则放大 FP32 权重更新 ⚠️ 注意 1. 缩放/反缩放必须是数学逆操作 2. optimizer.step 前必须反缩放 3. 动态S一般在2^8~2^16之间

看到这个流程了吗？缩放因子 S 就像一个桥梁，把 FP32 的损失和 FP16 的梯度连接起来。你要做的是让这座桥足够宽，保证梯度不会掉下去；同时又不能太宽，否则梯度直接溢出。

最后叮嘱一句：

我见过太多人上来就抄一个混合精度训练的代码，跑了几天发现效果不对，才回头补数学原理。别走这个弯路。花 30 分钟把这个流程理清楚，后面所有调优都能变成有方向的排查，而不是瞎试参数。

← 上一章 📖 返回目录 下一章 →