混合精度训练如何配置

混合精度训练（Mixed Precision Training）是一种通过同时使用 FP16（半精度）和 FP32（单精度）来加速训练、降低显存占用的技术，核心是利用 TensorCore（如 NVIDIA GPU）或专用硬件加速低精度计算。

下面按通用原理 → 主流框架配置 → 注意事项来说明。

一、混合精度训练的核心原理

1. 为什么要用混合精度

FP16：计算快、显存占用小（是 FP32 的一半）
FP32：数值稳定，避免下溢 / 梯度消失

2. 关键机制

权重主副本（Master Weights）：

优化器更新时使用 FP32 权重，前向/反向传播使用 FP16

Loss Scaling：

放大 loss，避免梯度在 FP16 下变成 0

二、PyTorch 中的混合精度配置（最常用）

✅ 推荐方式：`torch.cuda.amp`（自动混合精度）

1. 基本配置（PyTorch ≥ 1.6）

from torch.cuda.amp import autocast, GradScaler
import torch
scaler = GradScaler()

2. 训练循环示例

for data, target in dataloader:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
# 前向传播使用 autocast
with autocast():
output = model(data)
loss = criterion(output, target)
# 缩放 loss 并反向传播
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 不需要手动管理哪些层用 FP16 / FP32

3. 常见参数配置

参数	说明
`autocast()`	自动选择 FP16 / FP32
`GradScaler()`	防止梯度下溢
`scaler.scale()`	放大 loss
`scaler.step()`	优化器更新
`scaler.update()`	动态调整 scale

4. 多 GPU（DDP）

model = torch.nn.parallel.DistributedDataParallel(model)

✅ `torch.cuda.amp` 与 DDP 完全兼容

三、TensorFlow / Keras 配置

✅ TensorFlow 2.x（推荐）

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

模型示例

model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation="relu"),
tf.keras.layers.Dense(10, dtype="float32")  # 输出层建议 FP32
])

✅ 输出层、loss 计算建议保持 FP32

四、DeepSpeed / HuggingFace 配置（大模型常用）

1. DeepSpeed 混合精度

{
"fp16": {
"enabled": true
}
}

或 BF16：

{
"bf16": {
"enabled": true
}
}

2. HuggingFace Trainer

args = TrainingArguments(
fp16=True
)

或：

args = TrainingArguments(
bf16=True
)

五、FP16 vs BF16 如何选择？

类型	优点	缺点	适用
FP16	快、省显存	易溢出	NVIDIA GPU（TensorCore）
BF16	稳定、范围大	稍慢	A100 / TPU / 大模型

✅ 大模型 / 长时间训练 → BF16 更稳定

六、常见坑 & 注意事项

⚠️ 1. BatchNorm

通常保持 FP32
PyTorch autocast 已自动处理

⚠️ 2. Loss 数值异常

检查是否使用 GradScaler
不要手动 cast loss

⚠️ 3. 自定义 CUDA 算子

可能不支持 FP16
可强制 autocast(enabled=False)

七、快速检查是否生效

print(torch.cuda.get_device_properties(0))
# 看是否使用了 TensorCore（FP16）

或：

nvidia-smi
# 显存占用是否明显下降

八、什么时候该用混合精度？

✅ 推荐：

GPU 训练（尤其是 A100 / V100 / 3090）
大模型
显存不够
训练速度慢

❌ 不推荐：

CPU 训练
极大数值范围的模型

如果你愿意，可以告诉我：

用的 PyTorch / TensorFlow / HuggingFace
是否 多卡 / 大模型
是否有 OOM 或 NaN 问题

我可以直接给你完整可用的训练脚本配置。