混合精度训练怎么配置

混合精度训练（Mixed Precision Training）通过用 FP16/BF16 存模型和梯度、用 FP32 保精度，来加速训练、省显存。下面按框架给你一份“能直接照着配”的清单。

一、PyTorch（最常用）

1️⃣ 原生方案：torch.cuda.amp（推荐）

适合 PyTorch ≥ 1.6

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

✅ 关键配置点：

• autocast：自动选层用 FP16
• GradScaler：防止 FP16 梯度下溢
• BF16 不需要 GradScaler（A100 / H100 推荐）

2️⃣ PyTorch + Accelerate（最省心）

accelerate config
accelerate launch train.py

配置里选：

• Mixed Precision → fp16 / bf16
• DeepSpeed / FSDP 可选

3️⃣ PyTorch + DeepSpeed

{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000
}
}

或 BF16：

{
"bf16": { "enabled": "auto" }
}

二、Hugging Face Transformers

方法 1：Trainer

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

方法 2：命令行

--fp16
--bf16

三、TensorFlow / Keras

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

⚠️ 注意：

• 最后一层用 float32
• 优化器通常自动处理

四、关键配置建议（很重要）

✅ 显存 & 加速

✅ 常见坑

• ❌ 没用 GradScaler（FP16）
• ❌ loss 变成 NaN → 降低学习率 / 用 BF16
• ❌ 自定义 CUDA op 不支持 half

五、什么时候不该用？

• 数值极敏感（如 PINN、科学计算）
• 小模型 + 小 batch（加速不明显）

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

• 用的 框架
• 模型大小 / GPU
• 是否 多卡 / 分布式

我可以给你一份完整可直接跑的配置。