3、数据类型优化：FP32 vs FP16 vs BF16、混合精度训练原理、AMP自动混合精度实战、常见精度问题排查
好，咱们直接进入正题。数据类型优化，说白了就是怎么用更少的内存去存模型参数和中间结果。我刚开始做训练的时候，一张 24G 的卡，跑个 Bert-Large 都费劲。后来发现，其实很多地方根本不需要那么高的精度。
3.1 FP32、FP16、BF16 到底有什么区别？
先看一张表，把三种主流数据类型摆在一起对比一下：
| 数据类型 | 位宽 | 指数位 | 尾数位 | 表示范围 | 精度 |
|---|---|---|---|---|---|
| FP32 | 32 bit | 8 bit | 23 bit | ~3.4e38 | 高（约7位有效数字） |
| FP16 | 16 bit | 5 bit | 10 bit | ~65504 | 中（约3位有效数字） |
| BF16 | 16 bit | 8 bit | 7 bit | ~3.4e38 | 低（约2位有效数字） |
FP32 是默认精度，稳，但占内存。FP16 内存减半，但范围小，容易溢出。BF16 是 Google 搞出来的，保留了 FP32 的指数范围，牺牲了尾数精度。说白了，BF16 不怕溢出，但小数精度差点。
我在项目中遇到过一个问题：用 FP16 训练一个 7B 的模型，loss 直接变成 NaN。查了半天，发现是梯度更新时数值溢出了。换成 BF16 就没事了。嗯，这里要注意，不是所有硬件都支持 BF16，A100 及以上才原生支持。
3.2 混合精度训练的原理
混合精度训练，核心思路就一句话：用 FP16/BF16 做前向和反向，用 FP32 做参数更新和累加。
为什么会这样？因为梯度累加时，FP16 的精度不够，多个小梯度加起来会丢失信息。而参数更新时，学习率通常很小（比如 1e-5），FP16 根本表示不了这么小的数。
我画了一张流程图，帮你理清整个流程：
你看，整个流程里，FP32 就像个「保险柜」，专门存最重要的东西——主权重和 Loss。而 FP16/BF16 负责干活，速度快、省内存。
核心原则：计算用低精度，存储用高精度。低精度负责「跑得快」，高精度负责「跑得准」。
3.3 AMP 自动混合精度实战
PyTorch 的 AMP（Automatic Mixed Precision）帮我们自动做了上面那套流程。你只需要加几行代码就行。
我个人习惯用 torch.cuda.amp 这个模块。来看一个完整的训练循环示例：
import torch
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()
就这几行，显存占用直接减半。我在项目中遇到过最夸张的一次，一个 13B 的模型，用 FP32 根本放不进单卡，换成 AMP 后，居然跑起来了。
小技巧：如果训练过程中 loss 突然变成 NaN，可以尝试增大 init_scale 参数（默认 2^16）。或者检查一下模型里有没有不稳定的操作，比如 softmax 的输入过大。
3.4 常见精度问题排查
用混合精度训练，最怕的就是精度问题。我总结了几种常见情况：
- Loss 变成 NaN：通常是梯度溢出。解决办法：增大 Loss Scaling 的初始值，或者换成 BF16。
- Loss 不下降：可能是精度丢失太多。检查一下模型里有没有
LayerNorm或BatchNorm，这些层建议用 FP32 计算。 - 验证集精度下降：可能是某些层对精度敏感。我建议对关键层手动指定精度：
# 对特定层使用 FP32
with autocast(enabled=False):
x = self.critical_layer(x)
警告：千万不要在 autocast 上下文里做 torch.where 或 torch.clamp 这类操作，它们可能会破坏梯度流。我曾经因为这个排查了整整两天。
嗯，最后说一句。混合精度训练不是银弹。如果你的模型本身就很小（比如几百万参数），用 FP32 完全没问题。但如果你在搞大模型，AMP 就是标配了。
我曾经在一个 70B 模型上做过对比：FP32 需要 8 张 A100，AMP 只需要 4 张。省下来的卡，干点啥不好呢？