3、梯度累积（Gradient Accumulation）：原理与实现、Batch Size与显存的关系、实战代码示例
好，咱们今天聊聊梯度累积。这个话题，说白了就是「显存不够，技巧来凑」。我刚开始做大模型训练那会儿，经常被显存卡得死死的。明明想跑个大一点的 batch，结果 OOM 报错直接把我劝退。后来我学乖了——用梯度累积，效果还真不赖。
3.1 为什么需要梯度累积？
先问个问题：你训练模型时，是不是总想把 batch size 设大一点？
大 batch size 的好处很明显——梯度更稳定，收敛更平滑。但问题是，显存是有限的。你想想看，一个 batch 里每张图片都要过一遍网络，中间激活值全得存着。batch size 翻倍，显存占用也差不多翻倍。
我遇到过最尴尬的一次：明明就差那么 2GB 显存，就能跑 32 的 batch size。结果只能降到 16，训练速度直接慢了一倍。后来我用了梯度累积，batch size 设成 4，累积 8 步，效果跟 32 的 batch 几乎一样，显存却省了一大截。
核心思想：梯度累积就是把一个大 batch 拆成多个小 batch，分别计算梯度，然后累加起来，最后统一更新参数。这样你只需要一个小 batch 的显存，却能模拟大 batch 的训练效果。
3.2 Batch Size 与显存的关系
咱们来算一笔账。显存占用主要来自三块：
- 模型参数：固定大小，跟 batch size 无关
- 中间激活值：跟 batch size 成正比，这是大头
- 梯度：跟 batch size 成正比，但比激活值小
举个例子，假设你有个 1B 参数的模型，每个参数 4 字节（FP32）：
| 组件 | 显存占用 | 说明 |
|---|---|---|
| 模型参数 | 4 GB | 固定 |
| 优化器状态 | 8 GB | Adam 需要额外存一阶和二阶矩 |
| 梯度 | 4 GB | 跟参数一样大 |
| 激活值（batch=1） | ~2 GB | 取决于网络深度 |
| 激活值（batch=32） | ~64 GB | 直接炸了 |
看到了吧？激活值才是真正的显存杀手。你用梯度累积，每次只算 batch=1 的激活值，显存瞬间从 64GB 降到 2GB。嗯，这就是它的魅力所在。
3.3 梯度累积的原理
原理其实很简单，就三步：
- 前向传播：用小 batch 算一遍，得到 loss
- 反向传播：算梯度，但先不更新参数，把梯度累加到一个缓冲区里
- 重复：重复 N 步，梯度累加了 N 次
- 更新：用累加后的梯度除以 N（取平均），然后更新参数
说白了，就是把 N 个小 batch 的梯度「攒」起来，凑成一个虚拟的大 batch。我习惯把 N 叫做「累积步数」，一般设成 4、8 或者 16。
小技巧：累积步数 N 乘以小 batch size，应该等于你目标的大 batch size。比如你想跑 batch=64，但显存只够跑 batch=8，那 N 就设成 8。
3.4 知识体系流程图
下面这张图，是我自己画的一个梯度累积的流程。你看一眼就明白了：
这张图里，红色虚线表示「还没攒够，继续算」，绿色实线表示「攒够了，更新参数」。逻辑很清晰吧？
3.5 实战代码示例
好了，理论说完了，咱们直接上代码。这是我用 PyTorch 写的一个梯度累积的模板，我在好几个项目里都用过这个结构：
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
# 模型、数据、优化器
model = MyModel().cuda()
dataloader = DataLoader(dataset, batch_size=4) # 小 batch
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# 梯度累积参数
accumulation_steps = 8 # 累积 8 步
optimizer.zero_grad() # 只清零一次
for epoch in range(num_epochs):
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.cuda(), labels.cuda()
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播，累积梯度
loss = loss / accumulation_steps # 重要！取平均
loss.backward()
# 每 accumulation_steps 步更新一次
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad() # 清零梯度
# 可选：打印日志
print(f'Epoch {epoch}, Step {i+1}, Loss: {loss.item():.4f}')注意：代码里有个关键细节——loss = loss / accumulation_steps。为什么要除？因为每个小 batch 的 loss 是独立的，如果不除，累积 N 步后的梯度就相当于 batch size 放大了 N 倍，学习率也得跟着调。我一开始就踩过这个坑，梯度直接炸了。
3.6 避坑指南
我做了这么多年 GPU 编程，梯度累积的坑我基本都踩过。给你列几个最常见的：
- Batch Normalization 的问题：BN 层在训练时会统计当前 batch 的均值和方差。如果你用小 batch 累积，每个小 batch 的统计量都不一样。我建议要么用 SyncBN，要么把 BN 换成 LayerNorm。
- 学习率调整：累积步数变了，等效 batch size 变了，学习率最好也调一下。我一般按比例缩放——batch size 翻倍，学习率也翻倍。
- 梯度裁剪：如果你用了梯度裁剪，记得在累积完再裁剪，而不是每个小 batch 都裁。否则梯度会被「削」得太狠。
我的习惯：每次改累积步数，我都会先跑一个小实验，对比一下 loss 曲线。如果曲线跟之前差不多，那就没问题。如果 loss 震荡得厉害，说明学习率或者累积步数没调好。
3.7 什么时候不该用梯度累积？
嗯，也不是所有场景都适合。我遇到过两种情况，梯度累积反而帮倒忙：
- 显存够用的时候：如果你能跑大 batch，就别折腾了。梯度累积会引入额外的同步开销，训练速度反而慢一点。
- 对实时性要求高：比如在线学习，每来一个样本就要更新一次。累积步数会引入延迟，不适合。
说白了，梯度累积就是个「显存不够时的折中方案」。它不能提升模型精度，也不能加速训练，它只是让你在有限的显存下，能跑更大的等效 batch size。
好了，这一章就到这儿。记住一句话：显存不够，累积来凑。但别滥用，合适才是最好的。