4、梯度累积与显存优化：梯度累积原理、激活检查点、混合精度训练
显存不够用，是训练大模型时最头疼的问题。我见过不少团队，模型设计得挺好，一跑就 OOM。其实很多时候不是模型太大，而是我们没用好手里的显存。这一章聊的三个技术——梯度累积、激活检查点、混合精度训练——就是最实用的三板斧。
核心思路：用时间换空间，用精度换空间，或者两者结合。说白了，就是让显存不够的时候，还能把模型跑起来。
4.1 梯度累积：小 batch 也能有大效果
先问一个问题：为什么我们总想用大 batch？因为梯度更稳，收敛更快。但大 batch 意味着更大的显存占用。显卡装不下怎么办？
梯度累积的思路很简单：把一个大 batch 拆成几个小 batch，分别前向和反向，梯度攒够了再更新参数。显存占用只跟小 batch 有关，但效果接近大 batch。
4.1.1 原理与实现
正常训练流程是：前向 → 反向 → 更新参数。梯度累积改成：前向 → 反向 → 累积梯度 → 重复 N 次 → 更新参数 → 清空梯度。
# 梯度累积伪代码
accumulation_steps = 4 # 累积 4 次再更新
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / accumulation_steps # 归一化
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
注意那个 loss / accumulation_steps。我刚开始写的时候忘了这步，结果 loss 直接大了 N 倍，模型训飞了。嗯，这个坑我踩过。
4.1.2 适用场景与注意事项
| 场景 | 推荐累积步数 | 说明 |
|---|---|---|
| 单卡显存不足 | 2~8 | 最常用，显存占用降为 1/N |
| 跨卡同步梯度 | 4~16 | 减少通信频率，提升吞吐 |
| 模拟超大 batch | 32~128 | 注意 BN 层的统计量会受影响 |
我曾经踩过的坑：梯度累积时，如果模型里有 BatchNorm，它的 running mean/var 会在每个小 batch 都更新。累积步数多了，统计量会偏。建议用 SyncBN 或者冻结 BN 的统计量。
4.2 激活检查点：用计算换显存
你想想看，训练时显存大头在哪？不是模型参数，是中间激活值。一个 1B 参数的模型，激活值可能吃掉好几倍的参数显存。反向传播需要这些激活值算梯度，但前向时全存着太浪费了。
激活检查点（Activation Checkpointing，也叫梯度检查点）的思路是：前向时不存所有激活值，只存几个关键节点。反向时再重新算一遍。典型的 trade-off：显存从 O(L) 降到 O(sqrt(L))，但多了一次前向计算。
4.2.1 怎么用
PyTorch 里用起来很简单，加一行代码就行：
import torch.utils.checkpoint as checkpoint
# 原始前向
# output = model(input)
# 使用激活检查点
output = checkpoint.checkpoint(model, input, use_reentrant=False)
但要注意，不是所有层都适合加检查点。我个人的习惯是：只对显存占用大的层做检查点，比如 Transformer 的 attention 和 FFN。太小的层加了反而浪费计算。
4.2.2 分段策略
实际项目中，我会把网络分成若干段，每段只存输入和输出：
# 分段检查点示例
def forward_with_checkpoint(model, x, num_segments=4):
layers = list(model.children())
seg_size = len(layers) // num_segments
for i in range(num_segments):
start = i * seg_size
end = (i + 1) * seg_size if i < num_segments - 1 else len(layers)
segment = nn.Sequential(*layers[start:end])
x = checkpoint.checkpoint(segment, x, use_reentrant=False)
return x
分段越多，显存省得越多，但计算量也越大。我一般设 4~8 段，显存能省 40%~60%，计算只多 20%~30%。
小技巧：如果模型里有 Dropout，检查点模式下要小心。因为重新前向时 dropout mask 会重新生成，导致梯度不一致。建议把 dropout 放在检查点段外面。
4.3 混合精度训练：半精度也能训好模型
混合精度训练（AMP）现在几乎是标配了。原理很简单：用 FP16 做前向和反向，用 FP32 维护权重副本。显存直接砍半，速度还能快 2~3 倍。
但为什么不是纯 FP16？因为梯度下溢。小梯度在 FP16 里直接变 0，模型训不动。所以需要三个关键设计：
- 权重副本：FP32 存一份，FP16 算一份，更新时用 FP32
- Loss Scaling：把 loss 放大，梯度也跟着放大，避免下溢
- 动态 Scale：自动调整放大倍数，稳定训练
4.3.1 PyTorch AMP 实战
PyTorch 1.6 之后自带 AMP，用起来非常顺手：
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for inputs, labels in dataloader:
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
就这么几行，显存省一半。我见过有人不用 AMP，硬上大 batch 导致 OOM，其实加个 scaler 就解决了。
4.3.2 精度与稳定性
| 精度模式 | 显存节省 | 速度提升 | 适用场景 |
|---|---|---|---|
| FP32 | 1x | 1x | 小模型、精度敏感任务 |
| AMP (FP16) | ~2x | 2~3x | 大部分训练场景 |
| AMP (BF16) | ~2x | 2~3x | A100/H100 推荐，更稳定 |
我曾经踩过的坑：AMP 下 loss 突然变成 NaN，多半是 scale 值太大或者模型里有不稳定的操作。检查一下有没有 torch.exp 或者 softmax 的输入范围异常。把 GradScaler 的 init_scale 调小一点，或者用 BF16 试试。
4.4 三合一：实战中的组合策略
这三个技术不是互斥的，可以一起用。我最近训一个 7B 模型，单卡 A100 80G，就这么配的：
- 梯度累积：累积 8 步，等效 batch size 64
- 激活检查点：每 4 层一个检查点，显存省 50%
- AMP：BF16 模式，显存再省一半
最终显存占用从 72G 降到 38G，还能塞下更大的 batch。你想想看，如果不做优化，7B 模型单卡根本跑不起来。
我的建议：先开 AMP，显存不够再加梯度累积，还不行就上激活检查点。这个顺序性价比最高，改动最小，效果最明显。
4.5 本章知识体系
下面这张图把三个技术的关系和 trade-off 画清楚了：
这三个技术，说白了就是用不同的资源换显存。梯度累积用时间，激活检查点用计算，混合精度用精度。实际项目中，我一般先上 AMP，再根据显存余量决定要不要加后面两个。记住一点：优化不是目的，让模型跑起来才是。