混合精度训练如何实现加速 - GPU

混合精度训练（Mixed Precision Training）通过同时使用 FP16（半精度）和 FP32（单精度）来加速模型训练，同时尽量保持数值稳定性。下面从原理、实现方式、硬件支持、关键技巧几个方面系统说明。

一、为什么混合精度能加速训练？

1. 计算更快

FP16 比 FP32 计算更省资源
现代 GPU（如 NVIDIA Volta、Ampere）有 Tensor Core
Tensor Core 在 FP16 上的吞吐量是 FP32 的数倍
矩阵乘法（GEMM）是训练瓶颈，收益最大

2. 显存占用更少

FP16 只占 2 字节，FP32 占 4 字节
好处：
能跑更大的 batch size
能训练更大的模型
减少显存碎片和 OOM

3. 传输更快

显存 ↔ 计算单元的数据搬运更少
对大模型尤其明显

二、混合精度训练的核心思想

不是全部用 FP16，而是“该用 FP16 用 FP16，该用 FP32 用 FP32”

关键原则

操作	精度
前向传播	FP16
激活值	FP16
梯度计算	FP16
权重更新	FP32
Batch Norm / Loss	FP32（或受控 FP16）

三、数值稳定性问题如何解决？

1. 精度损失的根源

FP16 表示范围小
小梯度（如 1e-8）会被舍入为 0
导致梯度消失或训练不稳定

2. 解决方案：权重拷贝（Master Weights）

核心机制：

维护一份 FP32 的权重副本
前向 / 反向：用 FP16
权重更新：用 FP32

流程如下：

FP32 weights
↓（cast）
FP16 weights → forward → loss → backward
↑
更新到 FP32 weights

四、Loss Scaling（损失缩放）

为什么需要？

反向传播时梯度很小
直接 FP16 会下溢（underflow）

做法

计算 loss 后 放大一个常数（如 1024、动态）
反向传播
梯度再 除以同一因子
再更新权重

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

✅ 现代框架（如 PyTorch AMP）已自动完成

五、主流实现方式

1. PyTorch（推荐）

自动混合精度（AMP）

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

✅ 自动：

FP16 / FP32 切换
Loss scaling
数值稳定

2. TensorFlow

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

或在自定义训练中手动控制。

3. DeepSpeed / Megatron / FSDP

大模型训练默认混合精度
支持：
BF16（比 FP16 更稳定）
ZeRO + AMP

六、哪些操作不适合 FP16？

操作	原因
Batch Norm	统计量不稳定
Softmax（极端值）	数值溢出
Loss 计算	精度敏感
优化器状态	推荐 FP32

框架会自动处理，无需手动干预

七、实测加速效果（经验值）

场景	加速比
CNN（ResNet）	1.5 – 2×
Transformer	1.3 – 1.8×
大模型（A100）	2 – 3×
显存节省	30% – 50%

八、BF16 vs FP16（补充说明）

	FP16	BF16
指数位	少	和 FP32 一样
稳定性	较差	更好
硬件要求	老 GPU	A100+
是否需要 loss scaling	是	否

✅ 新 GPU 优先用 BF16

九、一句话总结

混合精度训练通过 FP16 加速计算、减少显存，用 FP32 权重和 Loss Scaling 保持数值稳定，从而实现训练加速。
如果你愿意，我可以：

用 最小化代码示例 解释
对比 FP32 / FP16 / BF16 的数值差异
帮你判断 你当前模型是否适合混合精度