损失缩放技术

使用 $FP16$ 格式时面临的主要问题是其动态范围相较于 $FP32$ 有限。尽管 $FP16$ 提供了显著的内存和潜在的速度优势，但在反向传播期间计算的梯度很容易超出其可表示范围。小梯度可能变为零（下溢），从而丢失重要的更新信息；而大梯度可能变为无穷大或非数值（NaN）（溢出），导致训练过程崩溃。损失缩放是一项旨在缓解这些问题，特别是梯度下溢的重要技术。

其核心思想很直接：如果小梯度在 $FP16$ 中容易下溢，我们可以在反向传播期间将其转换为 $FP16$ 之前人为地放大它们。我们通过在启动反向传播之前，将计算出的损失值乘以一个缩放因子 $S$ 来实现这一点。

考虑标准的反向传播过程，其中梯度 $g$ 相对于损失 $L$ 计算：$g = \frac{\partial L}{\partial w}$。 通过损失缩放，我们计算相对于缩放后损失 $L_{scaled} = S \times L$ 的梯度 $g_{scaled}$： $g_{scaled} = \frac{\partial (S \times L)}{\partial w} = S \times \frac{\partial L}{\partial w} = S \times g$ 此缩放操作有效地提升了梯度值，使其在 $FP16$ 中表示时下溢的可能性降低。

当然，这些缩放后的梯度（$g_{scaled}$）不能直接被优化器使用，因为它们不代表原始损失的真实梯度。因此，在反向传播计算梯度（可能在 $FP16$ 中）之后，但在优化器更新模型权重（通常使用 $FP32$ 梯度）之前，我们必须通过将梯度除以相同的因子 $S$ 来“反缩放”它们： $g = \frac{g_{scaled}}{S}$ 这样就恢复了原始梯度大小，现在希望能避免 $FP16$ 下溢导致的信息丢失。整个过程在每次训练步进的训练循环中发生。

选择缩放因子 $S$ 的主要方法有两种：静态损失缩放和动态损失缩放。

静态损失缩放

这是一种更简单的方法。在训练开始时选择一个固定不变的缩放因子 $S$，并在整个训练过程中使用它。

# PyTorch 静态损失缩放示例

# 假设 S 是一个预先选定的常数缩放因子
S = 128.0
scaler = torch.cuda.amp.GradScaler(
    init_scale=S, growth_interval=100000000
) # 强制静态

optimizer.zero_grad()

# 使用自动混合精度进行前向传播
with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

# 手动缩放损失（或让 GradScaler 处理）
# scaled_loss = loss * S
# scaled_loss.backward()

# 使用 GradScaler 缩放损失并调用 backward
scaler.scale(loss).backward()

# 在优化器步进前反缩放梯度
# 注意：scaler.step 隐式处理反缩放
# 如果 scaler 发现非有限梯度，则跳过 optimizer.step
scaler.step(optimizer)

# 为下一次迭代更新缩放因子（静态模式下无操作）
scaler.update()

静态缩放的主要难点在于选择一个适当的 $S$ 值。

• 如果 $S$ 过小，可能不足以防止非常小的梯度下溢。
• 如果 $S$ 过大，虽然它可以防止最终梯度下溢，但可能导致反向传播链中的中间梯度（在混合精度中也使用 $FP16$ 算术计算）在反缩放步骤之前发生溢出。

寻找最佳静态 $S$ 值通常需要手动调整和实验，这可能非常耗时。如果训练过程中梯度大小发生显著变化，可能还需要进行调整。

动态损失缩放

动态损失缩放通过在训练期间自动调整缩放因子 $S$ 来解决静态方法的不足。典型算法如下：

1. 将 $S$ 初始化为一个相对较大的值。
2. 在每次 backward() 反向传播后，检查计算出的梯度（在反缩放之前）是否存在溢出（即是否存在 Inf 或 NaN 值）。
3. 如果检测到溢出：
   • 跳过此批次的优化器步进（以避免使用无效梯度损坏权重）。
   • 降低缩放因子 $S$（例如，除以 2）。
4. 如果连续若干步（growth_interval）未检测到溢出：
   • 增加缩放因子 $S$（例如，乘以 2）。这会试探是否可以使用更大的缩放因子，从而可能将更小的梯度也推入 $FP16$ 的可表示范围。

这种动态调整有助于保持尽可能大的缩放因子，且不会引起溢出，从而最大限度地保护防止下溢，而无需手动调整。

现代深度学习框架提供了自动处理此问题的工具。在 PyTorch 中，torch.cuda.amp.GradScaler 实现了动态损失缩放。

# PyTorch 使用 GradScaler 进行动态损失缩放示例

# 使用默认动态行为初始化 GradScaler
# growth_interval 决定了它尝试增加缩放因子的频率
scaler = torch.cuda.amp.GradScaler(
    init_scale=65536.0, growth_interval=2000
)

for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        optimizer.zero_grad()

        # 使用自动混合精度进行前向传播
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)

        # 缩放损失，对缩放后的损失调用 backward()
        # 以生成缩放后的梯度。
        scaler.scale(loss).backward()

        # scaler.step() 反缩放优化器参数所持有的梯度。
        # 如果梯度包含 Inf 或 NaN，则跳过 optimizer.step()。
        scaler.step(optimizer)

        # 为下一次迭代更新缩放因子。
        # 如果发现 Inf/NaN 梯度则降低缩放因子，
        # 如果达到增长间隔则增加。
        scaler.update()

        # ... 训练循环的其余部分 ...

使用 GradScaler 抽象了检查溢出和调整 $S$ 的复杂性。您只需如示例所示包装前向传播、损失缩放、反向传播和优化器步进。

与梯度裁剪的交互

梯度裁剪是另一种用于稳定训练的技术（在第 17 章中讨论），它经常与混合精度一起使用。重要的是要正确地将梯度裁剪与损失缩放结合起来。标准做法是：

1. 通过 scaler.scale(loss).backward() 计算缩放后的梯度（$g_{scaled}$）。
2. 首先反缩放梯度：scaler.unscale_(optimizer)。这会就地修改与优化器参数关联的梯度，将其恢复到原始比例，但现在是 $FP32$ 格式。
3. 对反缩放后的 $FP32$ 梯度执行梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)。
4. 执行优化器步进：scaler.step(optimizer)。请注意，如果已经调用过 scaler.unscale_，则 scaler.step 不会再次反缩放。
5. 更新缩放器：scaler.update()。

BF16 与损失缩放

回顾一下，bfloat16 ($BF16$) 格式与 $FP32$ 具有相同的动态范围，尽管精度降低。由于其范围比 $FP16$ 宽得多，因此在使用 $BF16$ 时，梯度溢出和下溢发生频率远不那么常见。因此，使用 BF16 混合精度训练时，损失缩放通常不必要。 这简化了相对于 $FP16$ 的训练设置，前提是您的硬件高效支持 $BF16$ 操作（如 NVIDIA Ampere 架构 GPU 和 Google TPU）。然而，仍然建议监控梯度，因为极端情况可能仍会从稳定技术中受益或需要它们。

总之，损失缩放，特别是如 GradScaler 等框架工具中实现的动态损失缩放，是使用 $FP16$ 格式进行稳定有效混合精度训练的必不可少的技术。它通过动态调整梯度大小来抵消 $FP16$ 有限的数值范围，从而在不牺牲训练稳定性的情况下，实现显著的内存节省和潜在的速度提升。