如何在混合精度训练中应用

混合精度训练（Mixed Precision Training）是一种通过同时使用 FP16（半精度）和 FP32（单精度）来加速训练、降低显存占用，同时几乎不损失模型精度的技术。下面从原理、实现方式、注意事项三个层面系统讲解如何在训练中应用。

一、为什么要用混合精度训练

1. 主要优势

• 训练速度更快：GPU（尤其是 TensorCore）在 FP16 上计算效率更高
• 显存占用更低：FP16 参数、激活、梯度占用显存减半
• 几乎不损失精度：通过 FP32 主权重和数值稳定策略保证

2. 基准对比

二、混合精度训练的核心思想

1. 权重副本（Master Weights）

• 前向 / 反向：使用 FP16
• 参数更新：使用 FP32 主权重

FP32 weights
↓ 转成 FP16
FP16 weights → forward → loss → backward → FP16 gradients
↓
FP32 gradients → 更新 FP32 weights

2. Loss Scaling（损失缩放）

FP16 的值范围小，梯度容易下溢（变成 0），因此：

• 在 loss 上乘以一个 scale 因子
• 反向传播后 再除回来

loss = loss * scale
gradients /= scale

三、在主流框架中如何应用

✅ PyTorch（推荐）

方法一：使用 torch.cuda.amp（最常用）

1️⃣ 基本示例

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, target in dataloader:
optimizer.zero_grad()
with autocast():  # 自动混合精度
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 几乎不需要改模型代码

2️⃣ 注意事项

• autocast() 只作用于：
• model()
• loss
• 不要在 autocast 外手动转 half()

方法二：手动 FP16（不推荐）

model = model.half()
input = input.half()

❌ 容易数值不稳定，调试困难

✅ TensorFlow / Keras

使用 mixed_float16

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy('mixed_float16')

或：

policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

✅ 自动处理 loss scaling

✅ 适合 CNN、Transformer

✅ Hugging Face Transformers

from transformers import Trainer, TrainingArguments
args = TrainingArguments(
...,
fp16=True             # 开启混合精度
)
trainer = Trainer(model, args=args)

四、哪些操作不适合 FP16

PyTorch autocast 会自动处理这些。

五、常见坑 & 解决方案

1️⃣ Loss 变成 NaN

✅ 原因：

• loss scale 太大
• 模型本身不稳定

✅ 解决：

scaler = GradScaler(init_scale=1024)

2️⃣ 显存没减少

✅ 检查：

• 是否真的用了 autocast
• 是否保存了 FP32 中间变量

3️⃣ 验证阶段也要用吗？

❌ 不需要

✅ 推理可用 FP16，但验证建议 FP32

六、什么时候该用混合精度

✅ 适合：

• GPU 支持 TensorCore（V100 / A100 / 3090 / 4090）
• 大模型、大 batch
• Transformer / CNN

❌ 不适合：

• 极小模型
• 数值极度敏感任务（部分科学计算）

七、一句话总结

混合精度训练 = FP16 计算 + FP32 主权重 + Loss Scaling
如果你愿意，我可以：

• 帮你把现有训练代码改成混合精度
• 对比 FP32 / FP16 / BF16 的差异
• 针对你使用的模型（Transformer / CNN / LLM）给出最佳实践