如何实现混合精度训练

混合精度训练（Mixed Precision Training）是一种在训练深度学习模型时，同时使用 FP16（半精度浮点数） 和 FP32（单精度浮点数） 的技术，目的是在 几乎不损失精度 的前提下，加快训练速度、减少显存占用。

下面从 原理 → 关键技术 → 实现方式（PyTorch / TensorFlow）→ 注意事项 系统讲解。

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

1. 传统训练的瓶颈

• FP32（32-bit float）
• 精度高，但：
• 显存占用大
• 计算慢
• 不适合大规模模型

2. 混合精度的优势

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

核心原则

前向传播和反向传播用 FP16，参数更新用 FP32

关键机制

1. FP16 计算

• 前向传播
• 梯度计算

1. FP32 主权重（Master Weights）

• 保存一份 FP32 的参数副本
• 用于参数更新，避免精度损失

1. Loss Scaling（损失缩放）

• 防止 FP16 下梯度下溢（变成 0）

三、关键技术点详解

1. Loss Scaling（损失缩放）

FP16 可表示的最小正数约为：

6e-5

梯度可能小于这个值 → 直接变成 0
✅ 解决方法：

• 在 反向传播前放大 loss
• 反向传播后再 缩小梯度

loss → loss × scale
backward()
grad → grad / scale

2. Master Weights（主权重）

流程：

FP32 weights
↓
copy to FP16
↓
forward / backward (FP16)
↓
update FP32 weights

避免：

• 梯度累积误差
• 长时间训练后模型退化

四、PyTorch 实现混合精度训练（推荐）

✅ 使用 torch.cuda.amp

1. 基本示例

import torch
from torch.cuda.amp import autocast, GradScaler
import torch.nn as nn
import torch.optim as optim
model = Model().cuda()
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()
for x, y in dataloader:
x, y = x.cuda(), y.cuda()
optimizer.zero_grad()
# 前向传播：自动使用 FP16
with autocast():
y_pred = model(x)
loss = criterion(y_pred, y)
# 反向传播 + 参数更新
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

2. 关键 API 说明

3. 分布式训练（DDP）

✅ 与 torch.nn.parallel.DistributedDataParallel 完全兼容

with autocast():
output = model(x)
loss = criterion(output, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

五、TensorFlow 实现混合精度

1. 开启混合精度

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

2. 示例

model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10, dtype='float32')  # 输出层用 FP32
])
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
model.fit(x_train, y_train, epochs=10)

✅ TensorFlow 自动处理：

• Loss scaling
• 变量类型转换

六、适用场景与限制

✅ 适合

• GPU 训练（A100 / V100 / 3090 / 4090）
• 大模型 / 大 batch
• Transformer、CNN、Diffusion 模型

❌ 不适合

• CPU 训练（无收益）
• 数值极其敏感的模型

七、常见问题与调试建议

1. Loss 变成 NaN

✅ 检查：

• Loss scaling 是否开启
• 学习率是否过大
• 数据是否异常

2. 精度下降

✅ 尝试：

• 降低学习率
• 使用动态 loss scaling
• 输出层保持 FP32

八、总结一句话

混合精度训练 = FP16 加速计算 + FP32 保证稳定 + Loss Scaling 防止下溢
如果你愿意，我可以：

• ✅ 帮你把现有训练代码改成混合精度
• ✅ 对比 FP32 / FP16 显存和速度
• ✅ 解释 AMP 在 Transformer / LLM 中的使用细节