什么是混合精度训练实现

混合精度训练（Mixed Precision Training）是一种在深度学习训练过程中，同时使用多种数值精度（通常是 FP16 半精度 + FP32 单精度）来加速训练、减少显存占用，同时尽量保持模型收敛效果和精度的技术。

下面从“是什么、为什么要做、怎么实现”三个层面说明。

一、什么是混合精度训练

在神经网络训练中：

• FP32（32 位浮点）：精度高，但计算慢、显存占用大
• FP16（16 位浮点）：计算快、显存小，但容易溢出 / 精度不足
• BF16（脑浮点）：动态范围大，更适合训练（近年常用）

混合精度训练 =

用低精度（FP16 / BF16）做大部分计算
用高精度（FP32）保存关键参数和梯度，防止精度丢失

二、为什么要使用混合精度

1. 训练更快

• GPU（如 NVIDIA Tensor Core）对 FP16 有数倍算力
• 矩阵乘法、卷积明显加速

2. 显存更省

• 激活值、梯度占用减半
• 可上更大 batch / 更大模型

3. 通信更快（多卡）

• 梯度同步数据量减半

三、混合精度训练的核心实现机制

1. 权重备份（Master Weights）

• 优化器更新使用 FP32 权重
• 前向/反向使用 FP16 权重
• 避免小数更新被“舍入掉”

FP32 weights ──copy──> FP16 weights
↑                   ↓
optimizer          forward/backward

2. Loss Scaling（损失缩放）

FP16 梯度太小会下溢为 0：

loss = loss * scale
loss.backward()
grad = grad / scale

• 放大 loss → 放大梯度
• 更新前再缩回
• 现代框架可自动动态缩放

3. 算子精度选择

• 矩阵乘、卷积：FP16
• Softmax、LayerNorm、Loss：FP32
• Reduce / Exp / Log：保持高精度

四、常见实现方式（框架级）

PyTorch（推荐）

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

TensorFlow / Keras

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

分布式（AMP + DDP / FSDP）

• 梯度仍是 FP16 通信
• 参数更新在 FP32 进行
• 大模型训练标配

五、BF16 vs FP16

现代训练（LLM）更推荐 BF16

六、一句话总结

混合精度训练实现 = 低精度计算 + 高精度存储 + 损失缩放 / 权重备份，用框架 AMP 自动完成。
如果你愿意，我也可以：

• 讲 LLM 中的混合精度
• 对比 AMP vs 手动实现
• 分析 为什么你的模型不收敛