如何做混合精度训练

混合精度训练（Mixed Precision Training）是一种在深度学习训练中同时使用 单精度浮点数（FP32） 和 半精度浮点数（FP16 或 BF16） 的技术，目的是：

• 减少显存占用
• 提高训练速度
• 在尽量不损失模型精度的前提下完成训练

下面从 原理 → 关键问题 → 常见实现方式 → 实战示例 讲清楚。

一、为什么需要混合精度

1. 精度对比

2. FP16 的问题

• 梯度下溢（underflow）：很小梯度变成 0
• 溢出（overflow）：激活值过大变成 inf
• 累加误差：矩阵乘精度下降

二、混合精度的核心思想

用 FP16 做“快的部分”，用 FP32 做“稳的部分”
通常包括：

• ✅ 前向传播：FP16
• ✅ 矩阵乘法 / 卷积：FP16
• ✅ 梯度计算：FP16
• ❗ 参数更新：FP32（master weights）
• ❗ Loss 缩放（Loss Scaling）

三、关键技术点

1. Master Weights（主权重）

• 模型保留一份 FP32 权重
• 每次训练：

1. 拷贝到 FP16
2. 前向 + 反向
3. 梯度转回 FP32
4. 更新 FP32 权重

2. Loss Scaling（损失缩放）

解决 FP16 梯度下溢：

loss *= scale
grad /= scale

• 静态 scaling（固定值）
• 动态 scaling（自动调整）

四、常见实现方式

✅ 1. PyTorch 原生（推荐）

（1）AMP（Automatic Mixed Precision）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
optimizer.zero_grad()
with autocast():
output = model(input)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 自动处理：

• 哪些层用 FP16
• Loss scaling
• 数值稳定性

✅ 2. Hugging Face Transformers

training_args = TrainingArguments(
fp16=True,   # 或 bf16=True
)

自动启用混合精度。

✅ 3. TensorFlow

tf.keras.mixed_precision.set_global_policy('mixed_float16')

五、BF16 vs FP16 怎么选？

✅ BF16 不需要 loss scaling，更稳定

六、常见坑 & 建议

⚠️ 注意事项

• BatchNorm 建议保持 FP32
• Softmax / Log 操作易溢出
• 学习率可能要微调
• 某些自定义 CUDA 算子不支持 FP16

✅ 建议

• 先在小数据集验证精度
• 使用 GradScaler
• 监控 inf / nan

七、总结一句话

混合精度训练 = FP16/BF16 加速 + FP32 保稳 + Loss Scaling
如果你愿意，我可以：

• 给你 完整训练脚本
• 讲 大模型混合精度（ZeRO + FP16）
• 对比 不同 GPU 的最佳配置

你现在是做 CV / NLP / 大模型？