怎样实现混合精度训练

混合精度训练（Mixed Precision Training）是指在训练深度学习模型时，同时使用 FP16（半精度）和 FP32（单精度），以加快训练速度、减少显存占用，同时尽量保持模型精度。

下面按「原理 → 实现方式 → 常见坑」系统说明。

一、核心原理

1. 为什么用混合精度

• FP16 优点
• 显存占用减半
• 计算更快（Tensor Core 支持）
• FP16 缺点
• 数值范围小，容易：
• 下溢（梯度变成 0）
• 溢出（loss 变 NaN）
• 累加误差大

2. 混合精度怎么解决

关键三步：

1. FP16 前向 & 反向：加速计算
2. FP32 权重副本（Master Weight）：保证更新稳定
3. Loss Scaling：防止梯度下溢

二、标准混合精度训练流程

✅ 训练步骤

1. 拷贝 FP32 模型权重
2. 前向：权重转 FP16
3. 计算 loss（FP16）
4. Loss Scale（乘一个常数，如 1024）
5. 反向：梯度为 FP16
6. 梯度反 Scale
7. 用 FP32 权重更新
8. 再把权重转回 FP16

三、常见实现方式

1️⃣ PyTorch（推荐）

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

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 / FP32 切换
• Loss Scaling
• 数值稳定

方法二：torch.nn.apex（旧方案）

from apex import amp
model, optimizer = amp.initialize(
model, optimizer, opt_level="O1"
)

⚠️ 已不推荐，优先用 amp

2️⃣ TensorFlow / Keras

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

TF 会自动：

• 使用 FP16 计算
• 保持 FP32 权重

3️⃣ Hugging Face Trainer

training_args = TrainingArguments(
fp16=True
)

四、典型超参数

五、常见问题 & 解决方案

❌ Loss 变成 NaN

✅ 检查：

• Loss scale 是否过大
• 是否有 softmax / log 不稳定
• 是否未使用 master weight

✅ 解决：

• 减小 loss scale
• 使用 clip_grad_norm_

❌ 精度下降明显

✅ 可能原因：

• 某些层不适合 FP16（如 LayerNorm）
• 数值敏感操作

✅ 解决：

• 白名单 FP32（PyTorch 自动处理）
• 使用 O1 而非 O3

六、适用场景

✅ 非常适合：

• Transformer
• CNN 大模型
• GPU（A100 / V100 / 3090）

❌ 不适合：

• 极小模型
• 数值极度敏感任务

七、一句话总结

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

• 给你 完整训练模板
• 针对 你的模型 / 框架 给具体代码
• 讲 AMP 内部机制

直接告诉我你的使用场景即可。