3、混合精度训练(AMP)：FP16与BF16原理、AMP自动混合精度实现、GradScaler梯度缩放、实战收益与陷阱。
聊到显存优化，混合精度训练绝对是个绕不开的话题。我刚开始接触这玩意儿的时候，心里也犯嘀咕：把精度砍一半，模型还能收敛吗？后来真香了。说白了，这就是用更少的显存，跑更大的模型，而且速度还能快上一截。
FP16与BF16：精度砍半，效果能行吗？
先说说FP16。它用16位来存一个数，相比FP32的32位，显存直接减半。但代价是数值范围窄了，精度也低了。我在项目中遇到过最典型的问题：梯度下溢。训练到一半，loss突然变成NaN，查了半天，原来是梯度太小，在FP16下直接变零了。
BF16呢？它是Google搞出来的。指数部分跟FP32一样，范围很宽，但尾数部分精度低。你想想看，这意味着什么？它不容易溢出，但精度确实不如FP16。我个人习惯是：如果模型对精度不敏感，比如CV里的分类任务，FP16就够用。如果是LLM这种大模型，BF16更稳。
核心区别一句话：FP16范围窄但精度高，BF16范围宽但精度低。选哪个，看你的模型怕不怕溢出。
| 类型 | 指数位 | 尾数位 | 数值范围 | 精度 |
|---|---|---|---|---|
| FP32 | 8 | 23 | ~3.4e38 | 高 |
| FP16 | 5 | 10 | ~65504 | 中 |
| BF16 | 8 | 7 | ~3.4e38 | 低 |
AMP自动混合精度：别自己瞎折腾
手动去调哪些层用FP16、哪些用FP32？太累了。PyTorch的AMP（Automatic Mixed Precision）就是干这个的。它会自动判断：前向传播用FP16加速，反向传播也用FP16，但权重更新和某些敏感操作（比如BatchNorm）保留FP32。
实现起来很简单，就几行代码：
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, label in dataloader:
optimizer.zero_grad()
with autocast():
output = model(data)
loss = loss_fn(output, label)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()嗯，这里要注意：autocast 只包裹前向传播和loss计算。反向传播和优化器更新，交给 scaler 处理。我曾经见过有人把整个训练循环都包进去，结果梯度爆炸了还不知道为什么。
GradScaler梯度缩放：解决下溢的杀手锏
为什么需要GradScaler？因为FP16能表示的最小正数大约是6e-8。如果梯度比这个还小，直接变零。那怎么办？在反向传播前，先把loss放大一个倍数（比如2^16），这样梯度也跟着放大。更新权重时再缩回来。
这就是 GradScaler 干的事。它会动态调整这个缩放因子：如果连续N步没有出现inf/NaN，就放大一点；如果出现了，就缩小。我习惯把初始缩放因子设成2^16，对于大部分模型都够用。
小技巧：如果训练过程中频繁出现梯度溢出，可以试试把 GradScaler 的 growth_interval 调大，或者干脆手动固定缩放因子。
实战收益与陷阱：别光看显存省了多少
收益很明显：显存占用直接砍半，训练速度提升20%-50%。我试过一个7B的LLM，FP32下显存爆了，换成BF16+AMP，不仅跑起来了，速度还快了30%。
但陷阱也不少。我列几个常见的：
- 梯度下溢：loss突然变NaN。解决办法：用BF16或者调大缩放因子。
- 精度损失：模型收敛变慢。解决办法：检查哪些层被强制转成了FP16，手动排除敏感层。
- 硬件不支持：老显卡（比如V100）对BF16支持不好。我建议先查一下显卡的算力版本。
避坑指南：我曾经在一个NLP任务上用了FP16，结果BLEU分数掉了2个点。排查了半天，发现是LayerNorm的精度不够。后来手动把LayerNorm层排除在AMP之外，问题解决。
知识体系图：AMP的核心逻辑
下面这张图帮你理清AMP的整体流程。从数据输入到权重更新，每一步的精度类型都标清楚了。
看到没？整个流程里，只有前向传播和反向传播用了低精度，其他关键步骤（数据输入、Loss计算、权重更新）都保留了FP32。这就是AMP的精髓：用最少的精度损失，换最大的显存收益。
最后说一句：别盲目追求FP16。如果你的模型训练不稳定，试试BF16，或者干脆把AMP关掉，用纯FP32跑一轮看看是不是精度问题。我踩过的坑，希望你别再踩一遍。