混合精度训练怎样配置

混合精度训练（Mixed Precision Training）可以在几乎不损失模型精度的前提下，显著减少显存占用、加快训练速度。下面从原理 → 常见框架配置 → 注意事项三个层面说明。

一、混合精度训练基本原理

混合精度通常使用：

• FP16（半精度）：用于前向、反向计算，提速、省显存
• FP32（单精度）：用于：
• 模型权重副本（master weights）
• 优化器状态
• 容易溢出的计算（如 loss scaling）

关键问题：

• FP16 动态范围小，容易溢出（NaN）
• 解决方式：Loss Scaling

二、不同框架的配置方法

1️⃣ PyTorch（推荐方式）

✅ 使用 torch.cuda.amp（最常用）

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

✅ 特点：

• 自动选择 FP16 / FP32
• 自动 loss scaling
• 支持 CNN / Transformer

2️⃣ PyTorch Lightning（最简单）

Trainer(
precision=16,
accelerator="gpu",
devices=1
)

✅ 自动处理：

• AMP
• Loss scaling
• 梯度处理

3️⃣ TensorFlow / Keras

✅ TF 2.x 原生支持

tf.config.optimizer.set_experimental_options(
{"auto_mixed_precision": True}
)
policy = tf.keras.mixed_precision.Policy("mixed_float16")
tf.keras.mixed_precision.set_global_policy(policy)

⚠️ 注意：

• 最后一层通常设为 float32
• Dense(softmax) 建议显式指定 dtype

4️⃣ DeepSpeed / Megatron（大模型）

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

或：

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

✅ 建议：

• A100 / H100 → BF16 更稳定
• V100 / T4 → FP16 + loss scaling

三、关键配置建议（非常重要）

✅ 1. 是否使用 BF16？

BF16 不需要 loss scaling，稳定性更好。

✅ 2. 常见坑

• ❌ 模型输出层用 FP16 → 改成 FP32
• ❌ 自定义 CUDA op 不支持 half
• ❌ 学习率过大导致 NaN
• ❌ BatchNorm 在 FP16 下不稳定（通常内部用 FP32）

✅ 3. 调试技巧

print(torch.cuda.amp.is_autocast_available())

出现 NaN：

• 降低学习率
• 调大 loss_scale_window
• 换 BF16

四、一句话总结

PyTorch 用 torch.cuda.amp，大模型用 DeepSpeed/BF16，A100 以上优先 BF16，V100 用 FP16 + loss scaling。
如果你愿意，可以告诉我：

• 使用的框架（PyTorch / TF / JAX）
• GPU 型号
• 模型类型（CNN / Transformer / LLM）

我可以给你精确可复制的配置代码。