使用 $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$ 来解决静态方法的不足。典型算法如下:将 $S$ 初始化为一个相对较大的值。在每次 backward() 反向传播后,检查计算出的梯度(在反缩放之前)是否存在溢出(即是否存在 Inf 或 NaN 值)。如果检测到溢出:跳过此批次的优化器步进(以避免使用无效梯度损坏权重)。降低缩放因子 $S$(例如,除以 2)。如果连续若干步(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 章中讨论),它经常与混合精度一起使用。重要的是要正确地将梯度裁剪与损失缩放结合起来。标准做法是:通过 scaler.scale(loss).backward() 计算缩放后的梯度($g_{scaled}$)。首先反缩放梯度:scaler.unscale_(optimizer)。这会就地修改与优化器参数关联的梯度,将其恢复到原始比例,但现在是 $FP32$ 格式。对反缩放后的 $FP32$ 梯度执行梯度裁剪:torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)。执行优化器步进:scaler.step(optimizer)。请注意,如果已经调用过 scaler.unscale_,则 scaler.step 不会再次反缩放。更新缩放器:scaler.update()。BF16 与损失缩放回顾一下,bfloat16 ($BF16$) 格式与 $FP32$ 具有相同的动态范围,尽管精度降低。由于其范围比 $FP16$ 宽得多,因此在使用 $BF16$ 时,梯度溢出和下溢发生频率远不那么常见。因此,使用 BF16 混合精度训练时,损失缩放通常不必要。 这简化了相对于 $FP16$ 的训练设置,前提是您的硬件高效支持 $BF16$ 操作(如 NVIDIA Ampere 架构 GPU 和 Google TPU)。然而,仍然建议监控梯度,因为极端情况可能仍会从稳定技术中受益或需要它们。总之,损失缩放,特别是如 GradScaler 等框架工具中实现的动态损失缩放,是使用 $FP16$ 格式进行稳定有效混合精度训练的必不可少的技术。它通过动态调整梯度大小来抵消 $FP16$ 有限的数值范围,从而在不牺牲训练稳定性的情况下,实现显著的内存节省和潜在的速度提升。