★ FEATURED ARTICLE 1. 显存告急时梯度累积到底在做什么显存不够用几乎是每个搞深度学习的人都会撞上的墙。你写好模型、调好数据管道满心期待地按下运行结果第一个 batch 还没跑完终端就甩出一行红字CUDA out of memory。这时候摆在面前的选择通常有三个换更小的模型、换更大的卡、或者想办法让现有的卡塞下更大的 batch。前两个方案一个伤效果、一个伤钱包而梯度累积恰好是第三条路里最实用、改动成本最低的一种。梯度累积的核心思想其实特别朴素。常规训练里我们每个 batch 算一次梯度、更新一次参数。梯度累积则是连续算好几个 batch 的梯度把它们加起来先不更新等攒够了再一次性更新。这样一来等效的 batch size 就等于“单次前向的 batch size × 累积步数”。显存占用只跟单次前向的 batch size 有关而优化效果却接近大 batch。换句话说它用时间换空间让你在小显存上跑出大 batch 的训练行为。这个技巧适合谁如果你手头只有一张消费级显卡想微调一个参数量不算小的模型或者你在做对比实验需要复现某个大 batch 的训练配置但硬件跟不上又或者你发现 batch size 太小导致 loss 震荡、训练不稳定那梯度累积就是你要找的东西。它不挑框架PyTorch、TensorFlow、JAX 都能实现原理一致差别只在写法。我见过不少人第一次听到梯度累积以为是什么高深技巧其实它的数学等价性非常干净。假设损失函数对每个样本的梯度是线性的那么把 N 个小 batch 的梯度平均和用一个大 batch 算出来的梯度在期望上是等价的。注意这里说的是期望实际会有细微差别后面讲 BatchNorm 的时候会专门说这个坑。理解了这一点你就明白为什么梯度累积能“骗过”优化器让它以为自己在处理一个大 batch。2. 从显存账本到等效 batch 的完整推导2.1 显存到底被谁吃掉了要理解梯度累积为什么省显存得先搞清楚训练时显存都花在哪。以 PyTorch 为例一次前向加反向的显存开销大致分四块模型参数、梯度、优化器状态、以及前向过程中保存的中间激活值。前两块相对固定优化器状态在 Adam 这类自适应方法里通常是参数量的两倍一阶矩和二阶矩而激活值才是真正的大头它跟 batch size 和序列长度成正比。这就是关键所在。当你把 batch size 从 8 提到 64参数量和优化器状态几乎不变但激活值会涨到原来的 8 倍。很多情况下正是激活值把显存撑爆的。梯度累积的做法是单次前向仍然只用 batch size 8激活值维持在小规模但连续做 8 次把每次算出的梯度累加到同一个梯度缓冲区里。梯度缓冲区的大小只跟参数量有关跟累积步数无关所以显存占用基本不变。用一个具体的账本来说明。假设模型有 1 亿参数用 Adam 优化器混合精度训练。参数本身约 200MBfp16梯度约 200MB优化器状态约 800MBfp32 的一阶二阶矩加主权重这些加起来约 1.2GB是固定开销。剩下的激活值batch size 为 8 时可能占 6GBbatch size 为 64 时可能占 48GB。一张 24GB 的卡batch size 8 能跑batch size 64 直接爆。梯度累积让你用 batch size 8 的激活开销达到 batch size 64 的等效效果代价是训练时间变成原来的 8 倍左右。2.2 等效 batch 的计算与学习率调整等效 batch size 的公式很简单effective_batch micro_batch × accum_steps × num_devices。这里的 micro_batch 是单次前向真正喂进去的样本数accum_steps 是累积多少步更新一次num_devices 是并行设备数。比如单卡、micro_batch 为 4、accum_steps 为 16等效 batch 就是 64。但这里有个容易被忽略的问题等效 batch 变了学习率要不要跟着变答案是要。大 batch 训练时梯度的方差更小、方向更稳通常可以承受更大的学习率。业界的经验法则是线性缩放batch size 扩大 k 倍学习率也扩大 k 倍。不过这个规则不是万能的当 batch 大到一定程度后收益会递减这时候要用平方根缩放或者带 warmup 的策略。我的建议是如果你只是想让训练跑起来、效果别掉太多先把学习率按线性缩放调一下然后观察前几百步的 loss 曲线。如果 loss 下降平稳说明缩放合理如果一开始就震荡甚至发散说明学习率偏大往回调一点或者加长 warmup 步数。梯度累积本身不改变优化器的更新频率但它改变了每次更新所依据的梯度质量所以学习率这块必须重新审视不能照搬小 batch 的超参。2.3 为什么不是简单地把 batch 调小有人会问既然显存不够我直接把 batch size 调小不就行了何必搞梯度累积这个问题问到点子上了。小 batch 训练确实能跑但会带来两个实际问题。第一梯度噪声大loss 曲线抖得厉害收敛过程不稳定尤其是 BatchNorm 层在小 batch 下统计量估计不准效果会明显变差。第二很多模型的训练配方是围绕特定 batch size 调出来的你贸然改小等于偏离了原配方复现结果对不上。梯度累积的价值就在于它让你在显存受限的前提下尽量贴近目标 batch size 的训练行为。它不是简单地缩小 batch而是把大 batch 拆成多次小计算再合并既控制了显存峰值又保留了梯度质量。对于那些对 batch size 敏感的模型比如对比学习、目标检测里的一些结构这个区别非常关键。3. PyTorch 里的落地写法与关键细节3.1 最小可用代码骨架先给一个能直接抄的骨架这是最朴素的实现方式不依赖任何额外库accum_steps 8 optimizer.zero_grad() for step, (inputs, targets) in enumerate(dataloader): outputs model(inputs) loss criterion(outputs, targets) loss loss / accum_steps loss.backward() if (step 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()这段代码里有三个细节必须说清楚。第一loss要除以accum_steps。因为 PyTorch 的backward()默认是累加梯度如果你不除累积后的梯度就是真实梯度的accum_steps倍相当于学习率被放大了。除以累积步数之后累加得到的梯度就是平均值和用大 batch 算出来的梯度在尺度上一致。第二optimizer.zero_grad()的位置。它必须在optimizer.step()之后、下一轮累积开始之前调用。如果你习惯在每个 iteration 开头清零那要改成条件清零否则梯度刚攒起来就被清掉了训练完全失效。这是新手最常踩的坑没有之一。第三循环结束时可能有余数。如果数据总量不能被accum_steps整除最后几个 batch 的梯度会攒着没更新。稳妥的做法是在循环结束后判断一下如果还有未更新的梯度补一次optimizer.step()。3.2 混合精度下的梯度累积现在训练基本都开混合精度梯度累积和 AMP 配合时要注意 scaler 的处理。用torch.cuda.amp的话scaler.scale(loss).backward()会把 loss 放大后再反传scaler.step(optimizer)内部会做梯度裁剪和 unscale。累积场景下正确的写法是每个 micro step 都调用scaler.scale(loss).backward()但只在真正更新那一步调用scaler.step(optimizer)和scaler.update()。scaler torch.cuda.amp.GradScaler() for step, (inputs, targets) in enumerate(dataloader): with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, targets) / accum_steps scaler.scale(loss).backward() if (step 1) % accum_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()这里有个隐藏问题scaler.update()会根据是否发生梯度溢出动态调整缩放因子。在累积模式下它只在更新步被调用频率变低了缩放因子的调整节奏也跟着变慢。实测下来这通常不是大问题但如果你发现训练前期 loss 异常可以适当增加init_scale或者把growth_interval调小一点让 scaler 更快适应。3.3 梯度裁剪该放在哪一步梯度裁剪和梯度累积的顺序很容易搞错。正确的做法是先等所有 micro step 的梯度累积完再对累积后的总梯度做裁剪然后才更新参数。如果你在每个 micro step 都裁剪一次那裁的是局部梯度和最终用于更新的梯度不是一回事效果会打折扣。if (step 1) % accum_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() optimizer.zero_grad()注意裁剪的阈值也要相应调整。因为累积后的梯度是平均梯度量级和单步梯度差不多所以max_norm不用改。但如果你忘了除以accum_steps梯度被放大了裁剪就会频繁触发等于变相把学习率压得很低训练会变得很慢。3.4 用 Accelerate 或 Lightning 省心实现手写梯度累积不难但容易在细节上翻车。如果你用 HuggingFace 的 Accelerate它内置了gradient_accumulation_steps参数配合accelerator.accumulate(model)上下文管理器能自动处理清零、缩放、混合精度和分布式同步省去很多心智负担。PyTorch Lightning 里则是在 Trainer 里设accumulate_grad_batches同样开箱即用。不过用封装库有个代价你对底层行为的控制变弱了。比如你想在累积过程中做自定义的梯度操作或者调试梯度数值封装层会挡住一部分细节。我的建议是先用裸写法跑通一遍把原理吃透再换封装库提效。这样出了问题你知道去哪找原因而不是对着黑盒干瞪眼。4. 那些文档不会告诉你的坑4.1 BatchNorm 是最大的隐患梯度累积在数学上等价于大 batch这个结论对绝大多数层成立但 BatchNorm 是例外。BatchNorm 在前向时用当前 micro batch 的均值和方差做归一化并更新 running statistics。当你用 micro batch 为 4 做累积时BatchNorm 看到的一直是 4 个样本的统计量而不是等效的 32 个样本。这会导致两个后果一是归一化用的统计量不准二是 running mean 和 running variance 的更新节奏和真实大 batch 不一致。这个问题在微调视觉模型时特别明显。我的处理办法有几个如果模型允许把 BatchNorm 换成 GroupNorm 或 LayerNorm它们不依赖 batch 维度天然免疫这个问题。如果必须保留 BatchNorm可以考虑用sync_batchnorm在累积步之间同步统计量但这需要改模型结构成本较高。最省事的办法是在累积步数较大时适当调大 BatchNorm 的 momentum让 running statistics 更新得更平滑一些缓解统计量偏差。4.2 学习率调度器的步数会变学习率调度器通常是按 optimizer 的更新步数来走的比如 CosineAnnealing 或 StepLR。用了梯度累积之后optimizer 的更新频率变成了原来的 1/accum_steps如果你还按原来的总步数配置调度器学习率会在训练还没结束时就降到很低模型欠拟合。解决办法是重新计算总更新步数total_steps 原总步数 / accum_steps。然后在配置调度器时用这个新值。或者更简单把调度器的 step 调用放在optimizer.step()旁边让它跟着更新步走而不是跟着数据迭代步走。这个细节看起来小但影响的是整个训练的学习率曲线搞错了效果差很多。4.3 数据加载器里的 shuffle 和 drop_last梯度累积对数据顺序比普通训练更敏感。因为一次参数更新依赖连续多个 batch 的梯度如果这些 batch 的样本分布差异很大累积出来的梯度方向就会偏。所以shuffleTrue基本是必须的让每个累积窗口内的样本尽量随机。另外drop_last建议设为 True避免最后一个不完整的 batch 破坏累积节奏。如果数据量本来就少不舍得丢那就得在循环末尾手动处理余数确保梯度被正确更新。4.4 分布式训练下的累积多卡训练时梯度累积和 DDP 的配合需要额外注意。DDP 在每次backward()时会自动做 all-reduce 同步梯度如果你在每个 micro step 都触发同步通信开销会成倍增加而且同步的是局部梯度逻辑上也不对。正确的做法是用model.no_sync()上下文在非更新步跳过梯度同步只在最后更新那一步做一次 all-reduce。for step, (inputs, targets) in enumerate(dataloader): with model.no_sync() if (step 1) % accum_steps ! 0 else contextlib.nullcontext(): outputs model(inputs) loss criterion(outputs, targets) / accum_steps loss.backward() if (step 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()这样通信量降到原来的 1/accum_steps训练速度提升明显。不过no_sync用多了要注意某些版本的 PyTorch 在长时间不同步后可能有数值问题实测中偶尔遇到梯度 bucket 状态异常重启训练即可不算大毛病。5. 常见问题速查与排查思路5.1 训练不收敛或 loss 异常遇到 loss 不降反升先检查三件事。第一loss有没有除以accum_steps忘了除等于学习率放大数倍必炸。第二zero_grad()的位置对不对如果每个 iteration 都清零梯度根本攒不起来模型等于没更新。第三学习率有没有按等效 batch 重新调沿用旧学习率可能偏小导致收敛慢也可能偏大导致震荡。排查时可以打印每次optimizer.step()前的梯度范数和不用累积时的梯度范数对比。如果量级差很多基本就是缩放或清零的问题。另外可以在小数据集上跑几十步看 loss 是否能稳定下降快速定位是配置问题还是数据问题。5.2 显存没降下来理论上梯度累积不增加显存但如果你发现显存占用和不用累积时差不多甚至更高可能是这几个原因。一是loss没有及时释放计算图被保留着检查有没有在循环里累积 loss 张量。二是数据加载器的num_workers太多每个 worker 都占一份内存和显存映射。三是模型本身有缓存机制比如某些 attention 实现会缓存中间结果。用torch.cuda.memory_summary()看一下显存分布通常能定位到元凶。5.3 训练速度慢得离谱梯度累积用时间换空间慢是正常的但慢到不可接受就要查。如果 accum_steps 是 8训练时间大约是原来的 8 倍这是预期内的。如果超过这个比例可能是no_sync没用对分布式下通信开销没省下来。也可能是数据加载成了瓶颈micro batch 太小导致 GPU 利用率低这时候可以适当增大 micro batch减少 accum_steps在显存允许范围内找平衡点。问题现象可能原因排查方法解决方向loss 不下降忘记除以 accum_steps打印梯度范数loss 除以累积步数梯度攒不起来zero_grad 位置错误检查清零调用位置改为更新后清零显存没降计算图未释放memory_summary 查看避免累积 loss 张量训练过慢分布式未跳过同步检查 no_sync 使用非更新步 no_sync效果差于大 batchBatchNorm 统计偏差对比 running stats换 Norm 或调 momentum5.4 累积步数怎么选accum_steps 不是越大越好。步数太大参数更新频率太低训练总时间线性增长而且梯度累积窗口内的样本多样性下降可能影响泛化。我的经验是先确定目标等效 batch然后用显存能承受的最大 micro batch 去除得到的商就是 accum_steps。比如目标 batch 是 128显存最多支持 micro batch 16那 accum_steps 就是 8。如果算出来 accum_steps 超过 16就要考虑是不是 micro batch 设得太保守或者目标 batch 定得太激进。另外累积步数最好和数据集大小、总训练步数匹配。如果总步数本来就少accum_steps 太大可能导致整个训练只更新几十次模型根本没学够。这种情况下要么增加 epoch 数要么降低目标 batch。6. 几个实战中的经验补充梯度累积这个技巧我用了好几年踩过的坑基本都写在上面了。最后再补几个零散但有用的点。第一调试阶段可以把 accum_steps 设为 1确认基础训练流程没问题再逐步调大这样出问题容易定位。第二保存 checkpoint 时注意optimizer 的状态是在更新步才变化的如果你在累积中途保存恢复后可能丢失部分累积梯度建议只在更新步之后保存。第三如果你用梯度累积做梯度惩罚或对抗训练要特别小心这些方法本身依赖梯度结构累积后的梯度语义可能和预期不同需要单独验证。还有一个容易被忽视的点梯度累积和梯度检查点gradient checkpointing可以叠加使用两者省显存的机制不同一个省激活值的时间维度一个省激活值的空间维度配合好了能把显存压到很低。但叠加之后训练速度会进一步下降适合显存极度紧张、对时间不敏感的场景。实际项目里我通常会把梯度累积和自动混合精度、梯度检查点一起作为“显存三件套”来用。先用混合精度砍一半显存再用梯度检查点砍激活值最后用梯度累积把 batch 补回来。这套组合下来一张中等显存的卡能跑的模型规模能上一个台阶。当然每加一层优化调试复杂度就高一分建议逐项开启、逐项验证别一次性全上出了问题不好拆。
阅读完成 · 觉得有帮助？
