PyTorch梯度累积实战：显存不够时的等效大batch训练技巧

显存爆掉的瞬间，我盯着屏幕上那行 CUDA out of memory 愣了好几秒。当时正在调一个TCN加Transformer的长序列预测模型，回看窗口拉到了256根K线，单条样本推下去显存就飙得厉害，batch size调到8都站不住。后来我把梯度累积接上，配合混合精度，同样的模型在等效batch 128的配置下稳跑，Loss曲线比之前用单张小batch硬train时平滑得多。这篇就把我在PyTorch里用梯度累积从入门到调优的完整记录写出来，包括原理、可抄的代码、实测提速经验和几个我踩到怀疑人生的坑。

1. 显存不够时，别急着调小batch size：先搞清楚梯度累积为什么有效

1.1 OOM的真正元凶：中间激活值

很多人一遇到显存溢出，第一反应是调小 batch_size。这确实是最快的手段，但它有代价。要理解代价，得先明白显存到底被谁吃掉了。模型参数本身通常不是大头，真正撑爆显存的是前向传播过程中保存下来的中间激活值——每一层卷积或者Transformer里每个attention头计算出来的特征图，都得暂存在显存里，供反向传播时算梯度用。

以Transformer为例，输入序列长度为L，batch为B，隐层维度为d，那么单层注意力的激活值规模正比于 B * L * L * heads。序列越长，这部分占用的显存呈平方级增长。我那个预测模型回看窗口256，稍一加大batch，激活值就把显存吃完了。如果你用nvidia-smi盯着显存看，会发现参数加优化器状态可能只占2~3GB，中间激活却能把剩下10GB全占满。

1.2 缩小batch_size的连锁反应

既然激活值跟batch成正比，那就把batch调小，8不行就4，4不行就2？这样做的直接后果是梯度噪声变大。梯度是多个样本梯度的平均，样本越少，平均值越容易受个别样本影响，训练曲线会抖得厉害。

更麻烦的是BatchNorm这类层。训练阶段BN依赖当前batch内的均值、方差做归一化，batch太小，统计量本身就不稳定，模型甚至会越训越差。我见过有人把batch压到2之后，loss变成一条锯齿线，怎么调学习率都没用，最后根本不是学习率的问题，是batch太小导致训练从根本上不稳定。

1.3 梯度累积的核心思想与适用场景

梯度累积的思路很直白：既然一个batch塞不下，那就把大batch拆成几个小的micro-batch，一个接一个地前向、反向，但是不急着更新参数，而是把每个micro-batch算出来的梯度都累加到参数缓冲区里，攒够几个micro-batch之后，再用这批累积梯度统一更新一次参数。

这样做的效果，相当于你用更少的显存，凑出了一个等效的大batch。假设本来想用batch 32，显存只允许micro-batch 8，那么跑4次前向/反向、累加梯度后再做一次优化器更新，等效batch就是32。

这套思路尤其适合长序列模型、大batch才能稳定收敛的任务，以及单卡显存不够但暂时加不了多卡的情况。工程上它不算多么高深，但用得好不好，差距非常大。

2. 从数学到代码：手写一个可用的梯度累积训练循环

2.1 梯度线性叠加：为什么累积N步等价于batch翻N倍

要理解梯度累积为什么在数学上站得住脚，先回忆一个基本性质：梯度算子是线性的。假设一个大batch由N个样本组成，损失函数是每个样本损失的平均值：

[
L_{total} = \frac{1}{N}\sum_{i=1}^{N} L_i
]

它的梯度就是：

[
\nabla L_{total} = \frac{1}{N}\sum_{i=1}^{N} \nabla L_i
]

如果把这N个样本拆成K个micro-batch，每个micro-batch包含m个样本（N = K × m），先算每个micro-batch的平均损失，关键一步是把这个loss除以K，那么第一个micro-batch的梯度是：

[
g_1 = \frac{1}{m}\sum_{j=1}^{m} \nabla L_j / K
]

把K个micro-batch的梯度都累加起来：

[
\sum_{t=1}^{K} \frac{1}{K} \cdot \frac{1}{m}\sum_{j \in batch_t} \nabla L_j = \frac{1}{Km}\sum_{i=1}^{N} \nabla L_i
]

恰好等于大batch的梯度。这就是为什么需要 loss / accumulation_steps：不除的话，等效的学习率会被放大K倍，收敛很可能直接发散。

2.2 最小实现：不需要改模型任何一行

PyTorch里梯度累积的实现非常简洁。默认情况下，loss.backward() 会把梯度累加到每个参数的 .grad 上，只有调用 optimizer.step() 才会更新参数。所以我们要做的，就是控制 optimizer.step() 和 optimizer.zero_grad() 的调用时机。

下面是基础版本，不涉及混合精度：

python复制model.train()
optimizer.zero_grad()

accumulation_steps = 4

for i, (inputs, labels) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps   # 关键：保持数学等价
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        # 可选：梯度裁剪
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()

这段代码可以直接替换你原来的训练循环。需要注意数据加载器 DataLoader 里的 batch_size 应该设为micro-batch的大小，比如8；accumulation_steps 设为4，等效batch就是32。

2.3 不要忘记除以accumulation_steps

我最开始写梯度累积时，只记得累积梯度，忘了对loss做除法，结果训练前几百步就loss爆炸。原因上面已经说过：累积K步后参数更新量会放大K倍，等于偷偷把学习率调大了K倍。

正确做法是在每次backward之前把loss除以累积步数：

python复制loss = criterion(outputs, labels) / accumulation_steps

这样每个micro-batch的梯度只贡献“大batch平均梯度”的一部分，累积K步之后正好等于完整的大batch梯度。

另外还有一个易错细节：如果你的模型有多个loss需要相加，比如 loss = loss_main + 0.1 * loss_aux，那要除以累积步数的是加完后的总loss，而不是每个分量分别除。否则总梯度会被多除一次。

3. 让"累积"真正快起来：AMP配合与三项实测提速经验

3.1 先纠正一个概念：累积本身不会加速计算，它降低的是显存峰值

说实话，标题里那个"超快"一开始让我很纠结——因为梯度累积的本质是牺牲时间换显存，它不会减少计算量，甚至因为多次前向/反向，总计算量几乎不变。那为什么某些场景下加了梯度累积反而训练更快？

原因在于：如果显存不够导致你只能跑batch 2，那么用上累积、配合混合精度后，你可能跑得动batch 8甚至16；同样一个epoch里，大步长更新带来的训练效率提升，会明显好过小batch反复震荡。 另外，如果原来因为OOM经常被迫重试、或者频繁触发显存碎片整理，那么把batch控制在一个合理范围后，整体吞吐也会提升。所以"超快"的真实含义是：在显存受限时，它让你能用上更大的等效batch，从而更快更稳地收敛。

3.2 GradScaler与梯度累积的正确搭配姿势

混合精度训练（AMP）在PyTorch里通常配合 torch.cuda.amp.GradScaler 使用，作用是防止fp16下梯度下溢。它跟梯度累积一起用的时候，顺序要比基础版严格得多。

基础的AMP版训练循环是这样：

python复制scaler = torch.cuda.amp.GradScaler()
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(train_loader):
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels) / accumulation_steps

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.unscale_(optimizer)                     # 先换回真实梯度
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)                         # 再step
        scaler.update()
        optimizer.zero_grad()

注意 scaler.step() 内部会处理梯度缩放，但如果我们想用 clip_grad_norm_，就必须先调用 scaler.unscale_(optimizer)。这一步是把之前缩放过的梯度还原成正常范围，否则裁剪阈值就白设了，而且这个 unscale_ 必须放在裁剪之前。顺序错一毫都可能让裁剪失效或梯度失真。

3.3 我给训练循环做的三个小改动，直观提升了吞吐

除了AMP本身带来的显存下降和速度提升，还有几个工程上的小改动，我实测下来对吞吐有可感知的帮助。

第一个是给 torch.backends.cudnn.benchmark 设成 True。它让cuDNN在卷积层自动搜索最优算法，虽然第一次迭代有额外开销，但对固定输入尺寸的训练来说，后续每一步都能省下不少时间。对CNN和很多含卷积的模型都适用；纯Transformer模型提升不大，但开了也无害。

第二个是数据加载配置。DataLoader 里把 num_workers 调高一些、pin_memory=True，让数据搬运不阻塞GPU计算。梯度累积以后，前向/反向是连续进行的，如果数据加载跟不上，GPU就会空转，训练速度立刻打回原形。我习惯根据CPU核数把 num_workers 设在4~8之间，再配合预取，效果很直观。

第三个是减少循环里的隐式设备同步。比如不要每步都打印loss到终端、不要在循环里频繁调用 .item() 之后做大量CPU端逻辑，这些都会触发GPU和CPU的同步，打破流水线。累积模式下一口气能跑好几个micro-batch，就更不该用无谓的同步把它切成一段段。如果想看loss曲线，攒几个step再统一打印或者写到TensorBoard，效果一样，速度快得多。

3.4 累积步数怎么选：等效batch、显存余量和收敛稳定性的三角关系

积累步数不是越大越好。它受三个因素制约：显存余量、收敛稳定性、训练时间。

从显存角度看，accumulation_steps 越大，等效batch越大，但显存占用并不会有太大变化，因为显存峰值取决于micro-batch大小。所以一个务实的策略是：先确定显存能够稳定承载的micro-batch size，再反推需要多少步才能达到目标等效batch。

举个例子，你的目标等效batch是128，显存实测micro-batch能跑16，那 accumulation_steps 就是 128 / 16 = 8。如果光看loss曲线觉得震荡太大，可以保守地把目标等效batch提到256，累积步数翻倍；如果发现训练收敛太慢，则说明等效batch可能偏保守了，可以适当上调micro-batch。

经验上，等效batch在几十到一两百之间，对大多数人来说是一个既能保证梯度稳定、又不会过度稀释梯度噪声的区间。太小体现不出累积的作用，太大则容易陷入尖锐极小值，泛化反而变差。

4. 梯度累积的五个真实深坑，每个我都替你踩过

4.1 BatchNorm在累积下会变得不稳定

这是梯度累积最常见的坑，也是很多人在视觉模型上跑出糟糕结果的原因。BatchNorm在训练时会用当前micro-batch内的均值和方差做归一化，并更新running stats。梯度累积并不会改变这一点——所以BN的统计量依然基于micro-batch的大小，而不是等效的大batch。

micro-batch较小时，BN估计的均值和方差噪声很大，训练不稳定，甚至在验证时出现训练和推理统计量不一致的问题。解决思路有三条：一是保证micro-batch不要过小，一般建议至少8或16，太小就要考虑换层；二是把BatchNorm换成GroupNorm或LayerNorm，它们不依赖batch维度，受累积影响小得多；三是如果一定要用BN且batch确实很小，可以让BN统计量也去做累积，但实现起来比较麻烦，一般没必要。

4.2 学习率要不要跟着总batch size一起调

按照linear scaling rule，batch size翻倍，学习率也可以翻倍。但在梯度累积场景里，这个规则不能机械套用。

原因在于，真实的大batch在计算BN统计量、Dropout采样等环节，都跟“拆开再累积”的等效batch不完全一致。尤其是Dropout，在每个micro-batch里都会重新采样，等效batch下的随机性和真实大batch并不等价。所以累积后直接按倍数提高学习率，风险很大。

我的经验是：在梯度累积启动时可以保持学习率不变，先让loss曲线平稳跑起来，再根据收敛速度微调。如果想激进一点，可以把学习率提高10%~30%而不是直接翻倍，配合warmup观察。总之，学习率调整的幅度要保守，优先保证稳定。

4.3 step/epoch/update三套计数器的语义混乱

梯度累积之后，“一个step”的语义变得非常暧昧。有人统计epoch，有人统计iteration，还有人统计optimizer更新次数，混在一起沟通特别容易出问题。

我建议代码里统一维护两个变量：global_step（数据迭代次数，也就是micro-batch的个数）和 global_update（参数更新次数）。每执行一次 optimizer.step()，global_update 加1。学习率调度器的设计要想清楚：PyTorch里很多scheduler是按 step() 次数来衰减的，如果你在每次micro-batch后都调用 scheduler.step()，那么学习率会被额外衰减了很多次。正确做法是只在执行 optimizer.step() 之后调用 scheduler.step()，或者用 OneCycleLR 这类按更新步数设计的调度器。

4.4 梯度裁剪和AMP累积的顺序，错一步就崩

这个坑在3.2节提过，但值得再单独强调一次：先 unscale_，再 clip_grad_norm_，最后 step。

如果不先unscale，clip_grad_norm_ 看到的梯度还是被缩放过的，裁剪阈值完全失真。极端情况下，梯度本身很小，unscale之前被裁剪掉，参数更新量就不对了。另一个相关坑是：如果某一个micro-batch的后向传播里出现NaN或Inf，scaler.step() 会跳过这轮更新，但梯度累积的缓冲没有被清空，下一轮累积会把脏梯度带进去，让后续训练一起崩。所以一旦发现 scaler 跳过step，最好主动调用 optimizer.zero_grad() 清一下缓冲，或者干脆在训练循环里加入梯度范数的合法性校验。

4.5 DDP下累积与梯度同步的配合

如果从单卡迁移到多卡（DistributedDataParallel），梯度累积需要额外注意。DDP默认情况下，每次 loss.backward() 都会触发梯度all-reduce，把各卡梯度做平均。如果累积K步，那就意味着K次额外的通信，虽然通信和计算有重叠，但通信开销依然不可忽视。

进阶做法是用 model.no_sync() 上下文管理器把前K-1个micro-batch包裹起来，跳过这些步的梯度同步，只在最后一个micro-batch结束后同步一次：

python复制for i, (inputs, labels) in enumerate(train_loader):
    is_last_step = (i + 1) % accumulation_steps == 0
    context = model.no_sync() if not is_last_step else torch.no_grad()
    with context:
        # 注意：这里不能用torch.no_grad()包backward
        pass

严格来说，DDP的no_sync需要跟backward配套，标准写法是：

python复制if (i + 1) % accumulation_steps == 0:
    loss.backward()          # 触发all-reduce
    optimizer.step()
    optimizer.zero_grad()
else:
    with model.no_sync():
        loss.backward()      # 不触发all-reduce，只算本地梯度

这样每个累积周期只做一次通信，多卡环境下能省下不少时间。

5. 实战复盘：用梯度累积训一个TCN+Transformer股票预测模型

5.1 这类模型为什么天生需要梯度累积

股票序列预测这类任务，输入通常是过去一段时间的多维度行情数据，比如开高低收、成交量、技术指标。我的模型里回看窗口是256根K线，每根K线32个特征，TCN负责提取局部形态，Transformer捕捉长程依赖。

Transformer的self-attention复杂度随序列长度平方增长，256长度的注意力矩阵已经不小了，再加上TCN的膨胀卷积激活值，显存占用非常可观。如果batch设得太大，一张24G的卡撑不住；设得太小，模型训练又不稳定。这种场景几乎是梯度累积的标准用户：模型大、序列长、显存紧。

5.2 一个可以直接抄作业的配置与训练循环

下面是我在实践里用下来的一个稳定配置：

训练循环核心代码如下：

python复制model.train()
optimizer.zero_grad()

for i, (features, targets) in enumerate(train_loader):
    features, targets = features.cuda(), targets.cuda()

    with torch.cuda.amp.autocast():
        preds = model(features)
        loss = criterion(preds, targets) / accumulation_steps

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

注意数据加载器里 batch_size=16，而不是128。每条样本都包含256个时间步、32维特征，直接喂128条进去，显存绝对爆掉。

5.3 实测效果与几个收尾建议

在开启梯度累积+AMP之后，训练过程中的显存占用从接近爆掉降到了显存总量的60%左右，等效batch却保持在了128。loss曲线明显比之前小batch训练时平滑，验证集上无论是回归还是方向准确率，都比小batch版本更稳定。

最后给三点收尾建议。

第一，评估和保存模型务必安排在完整累积周期结束后，也就是一次optimizer更新完成之后再执行，不要在mid-cycle的中间状态上做验证，那会是一个参数未更新的“半成品”模型，状态很迷惑。

第二，warmup很重要。梯度累积等效batch大，一开始就用较大学习率容易震荡，我会在开始阶段跑少量step做学习率预热，模型快速进入稳定区域。

第三，记录梯度范数。用TensorBoard或者直接打印梯度l2范数，能帮你快速判断累积和缩放是否有问题。正常情况下梯度范数应该平稳下降；如果出现NaN，优先怀疑AMP和累积的交互，而不是模型本身。

根据我个人经验，梯度累积是一项值得花半天时间彻底掌握的工程手段。它不能替代更大的显存，也不能降低计算量，但在“单卡、大模型、长序列、显存有限”这些现实约束下，它能用很小的改动换来训练稳定性和等效大batch，比单纯调小batch硬扛要靠谱得多。把这套用熟了，再往后切过去多卡分布式里遇到梯度同步、batch分配，你也会发现底层的思路都是一脉相承的。