混合精度训练怎么上手 - GPU

混合精度训练（Mixed Precision Training）核心思路是：用 FP16/BF16 做大部分计算，用 FP32 保存关键状态（如权重、优化器状态），从而在几乎不损失精度的前提下，加快训练速度、节省显存。

下面按「概念 → 上手步骤 → 常见框架示例 → 注意事项」给你一个可直接上手指南。

一、先搞清楚几个基础概念

1. 常见数值类型

FP32：32 位浮点，精度高，显存占用大
FP16：16 位浮点，速度快，但容易溢出/下溢
BF16：16 位，动态范围大，比 FP16 更稳定（A100 / 3090+ 支持好）

2. 混合精度的关键组件

Master weights（FP32 权重副本）
Loss Scaling（FP16 时非常重要）
CUDA / GPU 支持（Tensor Core）

二、最推荐的上手方式（按易到难）

✅ 方式 1：直接用框架自带 AMP（最简单）

强烈建议新手从这里开始

PyTorch（推荐）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
optimizer.zero_grad()
for x, y in dataloader:
x, y = x.cuda(), y.cuda()
with autocast(dtype=torch.float16):  # 或 torch.bfloat16
output = model(x)
loss = criterion(output, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 几乎不用改模型
✅ 自动处理 loss scaling

✅ 支持 CNN / Transformer

TensorFlow / Keras

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

然后正常写模型即可。

✅ 方式 2：使用 Hugging Face Trainer

如果你做 NLP / 大模型：

training_args = TrainingArguments(
fp16=True,        # 或 bf16=True
...
)

一行配置搞定。

✅ 方式 3：DeepSpeed / FSDP（大模型必学）

DeepSpeed ZeRO + FP16/BF16
支持千亿参数训练

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

三、FP16 vs BF16 怎么选？

场景	推荐
老 GPU（V100）	FP16
新 GPU（A100 / 4090）	BF16
训练不稳定	BF16
推理	FP16

四、常见坑 & 调试技巧

❗ 1. Loss 变成 NaN

FP16 时检查：
是否忘了 GradScaler
softmax / log 是否数值爆炸
解决：
换 BF16
降低学习率
对某些层强制 FP32

❗ 2. 某些操作不支持 FP16

例如：

softmax 大 logits
LayerNorm（建议 FP32）

PyTorch 会自动 fallback，但速度会下降。

❗ 3. 显存没明显下降？

检查：

优化器状态是否还是 FP32（正常）
embedding 是否很大
是否开了 gradient checkpointing

五、上手练习建议（按顺序）

✅ 用 CIFAR / MNIST 跑 PyTorch AMP
✅ 对比 FP32 vs FP16 的：

显存
速度
精度

✅ 换 BF16 看稳定性
✅ 上 Hugging Face + fp16=True
✅ 最后再碰 DeepSpeed

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

用的框架（PyTorch / TF / JAX）
GPU 型号
任务类型（CV / NLP / 大模型）

我可以直接给你可运行的训练模板。