混合精度训练显存优化实战 - 第3节

3. 损失缩放 (Loss Scaling)

混合精度训练里，最让人头疼的问题就是梯度下溢。说白了，就是数值太小，直接被 FP16 的精度给「吞掉」了。

我刚开始做混合精度训练时，就踩过这个坑。模型训练到一半，loss 突然不降了，梯度全是 0。查了半天才发现，不是模型有问题，是精度不够。

FP16 的「盲区」

FP16 能表示的最小正数是 2⁻²⁴，大约 6e-8。你想想看，如果某个梯度值是 1e-10，FP16 根本存不住，直接变成 0。

这就像用一把最小刻度是毫米的尺子，去量一根头发丝的直径。量不出来，只能报 0。

但 FP32 能表示的最小正数是 1.4e-45，差了 37 个数量级。所以 FP32 下不会出现这个问题。

损失缩放怎么解决？

思路其实很简单：在反向传播之前，先把 loss 放大一个倍数。这样计算出来的梯度也会跟着放大。等梯度更新权重时，再缩回来。

举个例子：

# 原始流程
loss = model(input)
loss.backward()  # 梯度可能很小，FP16 存不住

# 损失缩放流程
loss = model(input)
scaled_loss = loss * scale_factor  # 放大 loss
scaled_loss.backward()  # 梯度也被放大了
optimizer.step()  # 更新时自动缩回

嗯，这里要注意：放大倍数不能太大，否则梯度会溢出到 FP16 的最大值 65504。那就变成上溢了，同样有问题。

静态缩放 vs 动态缩放

损失缩放有两种策略：静态和动态。我个人的经验是，新手先用动态，老手可以试试静态。

静态损失缩放

固定一个缩放因子，从头用到尾。比如 128、256、1024 这些值。

优点：简单，没有额外计算开销。
缺点：你得手动调这个值。调小了，梯度还是下溢；调大了，梯度上溢。

我在项目中遇到过，用静态缩放时，同一个模型在不同数据集上，最优缩放因子能差 4 倍。每次换数据都要重新调，很烦。

缩放因子	效果	适用场景
128	保守，适合小模型	ResNet-50 等
1024	激进，适合大模型	BERT、 等
动态	自动调节	通用

动态损失缩放

动态缩放会自动调整缩放因子。核心逻辑是：

1. 如果连续 N 步都没有梯度溢出，就把缩放因子翻倍（比如乘以 2）
2. 如果某一步出现了梯度溢出，就跳过这一步，把缩放因子减半

为什么会这样设计？因为梯度溢出意味着缩放太大了，需要降下来。而长时间没溢出，说明还有放大空间，可以试试更大的缩放因子。

我曾经用动态缩放训练一个 6B 参数的模型，缩放因子从 256 自动涨到了 32768。如果让我手动调，我肯定不敢设这么大。

PyTorch 中的实现

PyTorch 从 1.6 开始，把混合精度训练封装到了 torch.cuda.amp 里。损失缩放也自动帮你处理好了。

基础用法

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # 默认是动态缩放

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # 缩放 loss 并反向传播
    scaler.scale(loss).backward()
    
    # 更新权重（内部会处理缩放）
    scaler.step(optimizer)
    
    # 更新缩放因子
    scaler.update()

你看，代码量很少。核心就三步：

1. scaler.scale(loss)：放大 loss
2. scaler.step(optimizer)：更新权重，内部自动缩回梯度
3. scaler.update()：根据梯度情况，调整缩放因子

手动控制缩放因子

如果你想用静态缩放，可以这样：

# 固定缩放因子为 1024
scaler = GradScaler(init_scale=1024, growth_interval=None)

把 growth_interval 设为 None，缩放因子就不会变了。

检查梯度溢出

有时候你想看看梯度有没有溢出，可以这样：

scaler.scale(loss).backward()

# 检查是否有溢出
if scaler._scale is not None:
    scale_val = scaler._scale.item()
    print(f"当前缩放因子: {scale_val}")

scaler.step(optimizer)
scaler.update()

知识体系图

下面这张图总结了损失缩放的核心逻辑：

避坑指南

最后分享几个我踩过的坑：

• 不要手动缩放梯度：我曾经试过在 backward 之后手动乘以缩放因子，结果和 GradScaler 的机制冲突了，梯度直接爆炸。用官方 API 就好。
• 动态缩放不要设 growth_interval 太小：默认是 2000 步。设成 100 的话，缩放因子会频繁变化，训练不稳定。
• 多 GPU 训练时注意：每个 GPU 的 GradScaler 是独立的。如果你用 DistributedDataParallel，记得在每张卡上都创建 scaler。

损失缩放是混合精度训练的「安全气囊」。没有它，FP16 训练基本跑不起来。但有了它，你只需要关注模型本身，精度问题交给 PyTorch 就好。