4、数据类型优化：使用fp16/bf16混合精度训练（AMP）
说到显存优化，数据类型这块绝对是性价比最高的切入点。我刚开始做大规模训练时，经常被显存爆掉搞得焦头烂额。后来发现，很多时候我们根本不需要用float32去存那些参数——说白了，精度够用就行。
4.1 为什么需要混合精度？
先看一组数据。同样是存一个数：
| 数据类型 | 占用字节 | 相比FP32节省 |
|---|---|---|
| FP32 | 4字节 | 0% |
| FP16 | 2字节 | 50% |
| BF16 | 2字节 | 50% |
你想想看，一个7B的模型，光参数就占28GB显存。如果换成FP16，直接降到14GB。这省出来的空间，够你塞多少batch size？
但问题来了——直接用FP16训练，模型容易崩。我在项目中遇到过，梯度下溢出、loss变成NaN，那叫一个头疼。所以就有了混合精度训练（AMP）：该用FP16的地方用FP16，该用FP32的地方用FP32。
核心思想：前向和反向用低精度加速，权重更新用高精度保稳。
4.2 FP16 vs BF16：怎么选？
这两个都是16位浮点数，但内部结构完全不同。我画了张图，你看一眼就明白了：
看到了吗？BF16的指数位和FP32一样多，所以它的数值范围和FP32几乎一致。这意味着什么？
- FP16：范围窄（±65504），容易溢出。我训练BERT-large时，loss经常直接变成NaN，后来发现是梯度太大了，FP16装不下。
- BF16：范围宽（同FP32），基本不会溢出。但尾数只有7位，精度稍差。不过对于深度学习来说，这点精度损失通常可以接受。
我的建议：如果你的显卡支持BF16（A100、H100、RTX 4090等），优先用BF16。不支持的话，FP16配合loss scaling也能用。
4.3 AMP的三种策略
PyTorch的AMP提供了三种策略，我一个个说：
策略一：自动混合精度（最推荐）
import torch
from torch.cuda.amp import autocast, GradScaler
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler() # 用于loss scaling
for data, target in dataloader:
optimizer.zero_grad()
with autocast(): # 自动选择精度
output = model(data)
loss = loss_fn(output, target)
scaler.scale(loss).backward() # 缩放loss防止下溢出
scaler.step(optimizer) # 反缩放梯度并更新
scaler.update() # 更新缩放因子这段代码我用了两年多，基本没出过问题。autocast会自动判断：卷积、线性层用FP16，BatchNorm、Softmax用FP32。你什么都不用管。
策略二：手动精度控制（灵活但麻烦）
with autocast(dtype=torch.float16): # 强制用FP16
output = model(data)
with autocast(dtype=torch.bfloat16): # 强制用BF16
output = model(data)我个人不太推荐这种方式。除非你很清楚每个算子需要什么精度，否则容易出问题。
策略三：完全低精度（不推荐）
model.half() # 所有参数转FP16
data = data.half()嗯，这里要注意。全模型转FP16，BatchNorm会崩。我见过有人这么干，训练到一半loss直接飞了。除非你的模型没有BatchNorm，否则别这么玩。
4.4 Loss Scaling：为什么需要它？
为什么会这样？FP16的精度范围是±65504，但梯度值可能小到1e-7。这么小的数，FP16根本存不了，直接变成0。这就是梯度下溢出。
Loss scaling的原理很简单：把loss放大N倍，反向传播时梯度也跟着放大N倍。这样梯度就不会下溢出了。更新权重时再缩小回来。
关键参数：init_scale=2**16（初始缩放因子），growth_factor=2.0（成功时翻倍），backoff_factor=0.5（溢出时减半）
GradScaler会自动调整缩放因子。如果连续N步没溢出，它就放大一点；如果溢出了，它就缩小。整个过程不需要你干预。
4.5 实战中的避坑指南
我曾经踩过的坑：
- 坑1：用了AMP但没加GradScaler。结果训练到一半loss变成NaN，我还以为是模型写错了，排查了两天。
- 坑2：在autocast外面手动调用了
loss.backward()。这样loss没有被缩放，梯度照样下溢出。 - 坑3：把BN层也转成了FP16。BN需要FP32的精度来统计均值和方差，转成FP16后模型直接不收敛。
这里给几个实用建议：
- 先跑一个batch验证：用AMP前，先用一个batch跑通，看看loss是否正常下降。
- 监控梯度统计：打印梯度的max/min/mean，如果出现大量0，说明下溢出了。
- BF16不需要loss scaling：因为BF16的范围和FP32一样，不会溢出。所以用BF16时，可以不用GradScaler。
- 注意分布式训练：DDP + AMP时，每个进程都要有自己的scaler。
4.6 性能实测数据
我在一个8卡A100的集群上做过测试，模型是LLaMA-7B：
| 配置 | 显存占用 | 训练速度 | 收敛效果 |
|---|---|---|---|
| FP32 | 28.5 GB | 1.0x (基准) | 基准 |
| AMP (FP16) | 16.2 GB | 1.8x | 99.7% |
| AMP (BF16) | 16.2 GB | 1.9x | 99.9% |
看到没？显存省了43%，速度翻了近一倍，精度几乎没损失。这就是AMP的魅力。
小技巧：如果你的模型对精度特别敏感（比如科学计算），可以在关键层手动指定FP32：self.critical_layer = nn.Linear(1024, 1024).float()。这样AMP会自动跳过这层。
好了，数据类型优化就讲到这里。记住一句话：能用低精度就用低精度，但别让精度成为瓶颈。AMP帮你做了99%的工作，你只需要加三行代码。