混合精度训练：FP16与BF16的实战之道

混合精度训练，说白了就是让模型跑得更快、显存更省。我刚开始接触这个技术时，心里也犯嘀咕：用半精度会不会掉精度？模型会不会不收敛？后来在几个大模型项目里摸爬滚打，才真正摸透了它的脾气。

FP16与BF16：两种半精度的恩怨

先说说FP16。它用16位表示一个数，1位符号位、5位指数位、10位尾数位。范围大概在 ±65504 之间。嗯，这个范围其实挺窄的。我在训练一个BERT-large模型时，梯度一涨就溢出了，损失直接变成NaN，气得我差点砸键盘。

BF16呢？它是Google搞出来的。1位符号位、8位指数位、7位尾数位。指数范围和FP32一样，都是 ±3.4×10³⁸。尾数精度差了点，但好处是不容易溢出。

特性	FP16	BF16
指数位	5位	8位
尾数位	10位	7位
表示范围	±65504	±3.4×10³⁸
精度	高	低
溢出风险	高	低

你想想看，FP16的指数范围只有5位，稍微大点的数就存不下。BF16虽然尾数精度低，但至少不会动不动就炸。我个人习惯，如果硬件支持BF16，优先用它。NVIDIA的A100、H100都原生支持BF16，效果确实好。

AMP实现机制：自动混合精度是怎么玩的

AMP的核心思路很简单：大部分计算用FP16/BF16，敏感部分保留FP32。具体怎么做？

1. 自动类型转换：前向传播时，自动把模型参数和输入转成半精度
2. 损失缩放：防止梯度下溢，用GradScaler放大损失
3. 权重更新：梯度转回FP32，更新FP32的权重副本

我记得第一次用PyTorch的AMP时，就三行代码的事：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

你看，就这么简单。autocast上下文管理器自动处理类型转换，GradScaler负责缩放梯度。我曾经在一个项目里忘了加scaler，结果训练到一半loss就崩了，排查了半天才发现是梯度下溢。

GradScaler使用：别小看这个缩放器

GradScaler是AMP的核心组件。它的工作原理是这样的：

• 前向传播时，把损失放大一个倍数（默认2^16）
• 反向传播后，梯度也相应放大
• 优化器更新前，把梯度缩小回去

为什么要这么做？因为FP16能表示的最小正数是 2^-24 ≈ 5.96×10⁻⁸。如果梯度比这个还小，就会变成0。放大损失后，梯度也跟着放大，就不会下溢了。

我曾经遇到过一个坑：训练一个生成模型时，loss一直在降，但生成效果越来越差。后来发现是GradScaler的缩放因子太大，导致梯度爆炸。嗯，这里要注意，scaler的scale值不是越大越好。

# 自定义缩放策略
scaler = GradScaler(init_scale=2**10)  # 初始缩放值小一点
# 或者动态调整
if scaler.get_scale() < 1.0:
    scaler.update(1.0)  # 防止缩放因子太小

实战对比：FP32 vs FP16 vs BF16

我拿一个ResNet-50在ImageNet上做了对比测试。硬件是A100，batch size=256。

精度模式	训练时间	显存占用	Top-1准确率
FP32	100%	100%	76.15%
FP16 AMP	62%	58%	76.12%
BF16 AMP	60%	57%	76.14%

结果很明显：混合精度训练快了近40%，显存省了40%多，准确率几乎没掉。BF16和FP16表现差不多，但BF16更稳定，没有出现梯度溢出的情况。

我个人建议，如果你的显卡支持BF16（比如A100、H100、RTX 3090以上），优先用BF16。如果不支持，FP16也够用，但记得加GradScaler。

这张图展示了混合精度训练的完整流程。从FP32权重开始，转成半精度做前向传播，计算损失并缩放，反向传播得到梯度，再反缩放回FP32更新权重。循环往复，直到训练结束。

最后说一句：混合精度训练不是银弹。如果你的模型本身就很小，或者对精度极其敏感（比如科学计算），那还是老老实实用FP32。但对于大多数深度学习任务，AMP是标配，不用反而奇怪。

想要更多资料，联系可可