混合精度训练有哪些实现步骤

混合精度训练（Mixed Precision Training）是一种通过结合半精度（FP16）和单精度（FP32）来加速深度学习训练、减少显存占用的技术。其核心思想是：大部分计算用 FP16，关键步骤保留 FP32 以保持数值稳定性。下面是通用实现步骤（以 PyTorch 为例，同时也适用于其他框架）。

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

• 前向传播 / 梯度计算：使用 FP16（更快、更省显存）
• 参数更新、Loss 缩放、Batch Norm：使用 FP32（防溢出 & 精度损失）

二、混合精度训练的主要实现步骤

✅ Step 1：启用混合精度上下文

使用 自动混合精度（AMP, Automatic Mixed Precision）。

PyTorch 示例

from torch.cuda.amp import autocast, GradScaler

✅ Step 2：创建 GradScaler（梯度缩放器）

作用：

• 防止 FP16 下梯度下溢（变成 0）
• 对 loss 进行放大，在反向传播后再缩小梯度

scaler = GradScaler()

✅ Step 3：修改训练循环（关键点）

1️⃣ 前向传播使用 autocast

for x, y in dataloader:
optimizer.zero_grad()
with autocast():  # FP16 自动转换
y_pred = model(x)
loss = criterion(y_pred, y)

✅ 在 autocast() 作用域内：

• 矩阵乘法、卷积等自动使用 FP16
• BN、Softmax 等计算仍可能保持 FP32

2️⃣ 使用 scaler.scale(loss) 进行反向传播

scaler.scale(loss).backward()

✅ 作用：

• 放大 loss，避免梯度下溢

3️⃣ 更新参数

scaler.step(optimizer)
scaler.update()

✅ scaler.step()：

• 自动处理梯度是否溢出
• 若溢出则跳过本次更新

✅ Step 4：保持 Master Weights（可选但推荐）

PyTorch AMP 自动维护 FP32 主权重：

• 模型参数：FP16（用于前向）
• 优化器状态：FP32（用于更新）

✅ 无需手动实现

✅ Step 5：Batch Normalization 注意事项

• BN 的统计量（running_mean / running_var）通常保持 FP32
• 在 autocast() 中 PyTorch 会自动处理

⚠️ 注意：

• 小 batch size + AMP 可能导致 BN 不稳定

✅ Step 6：验证 / 测试阶段

验证阶段 不需要 AMP：

model.eval()
with torch.no_grad():
for x, y in val_loader:
y_pred = model(x)

✅ 或显式关闭 autocast：

with autocast(enabled=False):
...

三、完整训练代码示例（PyTorch）

from torch.cuda.amp import autocast, GradScaler
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())
criterion = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(epochs):
model.train()
for x, y in dataloader:
x, y = x.cuda(), y.cuda()
optimizer.zero_grad()
with autocast():
y_pred = model(x)
loss = criterion(y_pred, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

四、TensorFlow / JAX 中的实现（简述）

TensorFlow

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

JAX

• 使用 jax.numpy.float16
• 手动管理 loss scaling（较复杂）

五、混合精度训练的优缺点

✅ 优点

• 训练速度提升（1.5x ~ 3x）
• 显存占用减少（约 30% ~ 50%）
• 对大模型尤其重要

⚠️ 缺点

• 数值稳定性问题（需 loss scaling）
• 某些算子不支持 FP16
• 精度可能轻微下降

六、什么时候不适合混合精度？

• 极小模型（收益不明显）
• 对数值精度极度敏感的任务
• 某些 RNN / 自定义算子

如果你愿意，我也可以：

• ✅ 对比 FP16 / FP32 / BF16
• ✅ 讲 Loss Scaling 的数学原理
• ✅ 结合 大模型（LLM）训练 讲混合精度
• ✅ 分析 AMP vs 手动混合精度

你更关心哪一部分？