混合精度训练 FP16 避坑指南：那个著名的“损失缩放 Loss Scaling”是干嘛的？

混合精度训练（Mixed Precision Training）是现代深度学习模型训练中常用的优化手段。通过将模型的大部分计算转移到半精度浮点数（FP16）进行，而保持关键部分（如权重更新）使用全精度浮点数（FP32），可以显著提高训练速度并减少显存占用。

然而，当迁移到FP16时，一个常见的陷阱是梯度下溢（Underflow）。为了解决这个问题，我们需要引入一个著名的技术：损失缩放（Loss Scaling）。

1. 为什么FP16会导致梯度下溢？

FP16的数值表示范围比FP32小得多。具体来说：

• FP32 (单精度): 最小可表示的正规数约为 $1.17 imes 10^{-38}$。
• FP16 (半精度): 最小可表示的正规数约为 $6.1 imes 10^{-5}$。

在深度学习的反向传播过程中，特别是对于深层网络，梯度值往往会变得非常小（例如 $10^{-6}$ 到 $10^{-8}$ 级别）。如果这些极小的梯度以FP16形式存储，它们将低于FP16的最小正规数，被自动“截断”为零（即下溢）。一旦梯度变成零，权重就无法更新，训练也就失败了。

2. 损失缩放（Loss Scaling）的工作原理

损失缩放的目的是在反向传播开始之前，人为地将这些极小的梯度值放大到一个FP16可以安全表示的范围内。

核心思想：

1. 放大损失： 在计算完损失 $L$ 后，将其乘以一个巨大的缩放因子 $S$（例如 $S=2^{10}=1024$ 或 $S=2^{15}=32768$），得到缩放后的损失 $L_{scaled} = L imes S$。
2. 反向传播： 对 $L_{scaled}$ 进行反向传播。根据链式法则，计算出的梯度 $
   abla W_{scaled}$ 也将被放大 $S$ 倍：
   $$\nabla W_{scaled} = \frac{\partial L_{scaled}}{\partial W} = S \times \frac{\partial L}{\partial W}$$
   现在，原本极小的 $\frac{\partial L}{\partial W}$ 经过 $S$ 倍放大后，就能安全地存储在FP16中。
3. 解除缩放： 在优化器（Optimizer）更新权重 $W$ 之前，我们需要将梯度除以 $S$ 还原，以确保权重更新的幅度正确。
   $$W \leftarrow W – \alpha \times \frac{\nabla W_{scaled}}{S}$$

3. PyTorch中的实操：使用 GradScaler

在PyTorch中，我们无需手动管理 $S$ 的选择和动态调整，torch.cuda.amp.GradScaler 负责实现动态损失缩放，这是使用PyTorch自动混合精度（AMP）的关键组件。

GradScaler 不仅执行缩放和解除缩放，它还会智能地监控梯度。如果梯度出现 Inf 或 NaN（通常是 $S$ 选得太大了，导致梯度溢出），它会自动降低 $S$；如果一段时间内没有出现问题，它会逐渐提高 $S$，以优化训练精度。

示例代码：

以下代码演示了如何在PyTorch中使用 GradScaler 进行混合精度训练：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast

# 确保CUDA可用
if not torch.cuda.is_available():
    raise EnvironmentError("CUDA is required for AMP training")
device = torch.device("cuda")

# 1. 初始化模型、优化器和损失函数
model = nn.Linear(10, 1).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# 2. 初始化GradScaler (损失缩放的核心)
scaler = GradScaler()
print(f"初始损失缩放因子: {scaler.get_scale():.1f}")

# 模拟数据
dummy_input = torch.randn(64, 10).to(device)
dummy_target = torch.randn(64, 1).to(device)

# 3. 训练迭代步骤
optimizer.zero_grad()

# 3a. 使用autocast上下文，自动将Tensor和操作转换为FP16
with autocast():
    output = model(dummy_input)
    loss = criterion(output, dummy_target)

# 3b. 使用 GradScaler 进行损失缩放并反向传播
# scaler.scale(loss) -> 损失 * 缩放因子S
scaler.scale(loss).backward()

# 3c. 梯度解除缩放并更新权重
# 1. 检查梯度是否有效（非Inf/NaN）。
# 2. 如果有效，解除缩放（梯度 / S），并调用 optimizer.step() 更新权重。
# 3. 如果无效，则跳过更新（保留原权重）。
scaler.step(optimizer)

# 3d. 更新缩放因子
# 根据上一步的检查结果，动态调整缩放因子S
scaler.update()

print(f"训练后的损失值: {loss.item():.6f}")
print(f"更新后的损失缩放因子: {scaler.get_scale():.1f}")

4. 避坑总结

• 必须使用 Loss Scaling： 在使用FP16训练时，如果禁用了损失缩放，几乎必然会导致梯度下溢，训练效果将大打折扣。
• 只缩放 Loss： 缩放操作只应用于损失值，不应用于激活值或其他计算结果。
• PyTorch推荐使用AMP： 推荐使用 torch.cuda.amp 模块进行自动混合精度，它集成了 autocast 和 GradScaler，极大简化了FP16的配置和管理工作。