混合精度训练为什么需要loss_scaler

混合精度训练需要loss_scaler的原因

混合精度训练中使用loss_scaler主要是为了解决梯度下溢(Gradient Underflow)问题。让我详细解释一下:

🎯 核心问题:FP16数值范围有限

FP16 vs FP32的表示范围对比

FP32(单精度): 
- 最小正数: ~1.4e-45
- 最大正数: ~3.4e+38

FP16(半精度):
- 最小正数: ~5.96e-8   ← 问题所在!
- 最大正数: ~65504

梯度下溢示例

# 训练过程中常见的梯度值
gradient_fp32 = 1e-10  # FP32可以正常表示

# 转换为FP16时
gradient_fp16 = float16(1e-10)  
# 结果: 0.0  ← 下溢为零!

# 后果: 参数无法更新
param = param - lr * gradient_fp16  
# 相当于 param = param - lr * 0 = param (没有变化)

🛠️ Loss Scaler的工作原理

三步救援机制

# ===== 步骤1: 前向传播时放大损失 =====
loss = criterion(outputs, labels)  # 假设loss = 1e-5
scaled_loss = loss * scale_factor  # scale_factor=1024
# scaled_loss = 0.01024  ← 现在FP16可以安全表示

# ===== 步骤2: 反向传播(梯度自动放大) =====
scaled_loss.backward()
# 链式法则: ∂(scaled_loss)/∂param = scale_factor × ∂loss/∂param
# 原始梯度 1e-10 → 放大后 1e-7  ← FP16可以表示!

# ===== 步骤3: 更新参数前缩小梯度 =====
for param in model.parameters():
    param.grad = param.grad / scale_factor  # 恢复原始梯度
    param.data = param.data - lr * param.grad  # 正常更新

📊 直观对比

场景	无Loss Scaler	有Loss Scaler
原始梯度	1e-10	1e-10
FP16存储	0.0 (下溢)	1e-7 × 1024 = 0.0001024
参数更新	❌ 无更新	✅ 正常更新