一、梯度累积：显存优化的核心武器

做深度学习训练的人，十有八九都遇到过显存不够的窘境。我当年第一次跑BERT-large时，看着那张"CUDA out of memory"的红字，真是欲哭无泪。后来发现，梯度累积这个技巧，说白了就是拿时间换空间——但换得值。

1.1 梯度累积的原理

先说说梯度累积到底在干什么。正常的训练流程是这样的：

前向传播算loss
反向传播算梯度
优化器更新参数

梯度累积的做法是：

前向传播算loss
反向传播算梯度（但不更新参数）
重复步骤1-2，累积梯度
累积到一定次数后，用累积的梯度更新参数

嗯，这里要注意：梯度累积不是把多个batch的数据拼在一起跑，而是分多次跑小batch，把梯度攒起来。我刚开始做的时候也搞混过，以为是把数据concat起来，结果显存直接炸了。

核心公式：

假设我们想要等效batch size为B，实际能承受的batch size为b，那么累积步数N = B / b。

参数更新时，实际使用的梯度 = 累积的梯度 / N

1.2 微批次(Micro-batch)设计

微批次就是每次实际送入GPU的那一小批数据。设计微批次时，我一般遵循几个原则：

微批次大小：以刚好填满GPU显存为准。我习惯先试一个batch size=1，看看显存占用，然后逐步往上加。
累积步数：通常设为2的幂次，比如2、4、8、16。为什么？因为GPU对2的幂次更友好，这是我踩坑踩出来的经验。
总batch size：微批次大小 × 累积步数。这个值要和你的学习率匹配。

举个例子，我去年做一个NLP项目，单卡只能跑batch size=4。但实验需要batch size=32的效果。怎么办？微批次设4，累积步数设8，完美解决。

1.3 与Batch Size的关系

这里有个容易翻车的地方：梯度累积不等于增大batch size。你想想看，真正的large batch是一次性看到更多数据，梯度更稳定。而梯度累积是分多次看到数据，梯度是逐步累积的。

我曾经犯过一个错误：把累积步数设得特别大（比如64），结果模型收敛特别慢。后来发现是学习率没调。一般来说：

等效Batch Size	学习率调整建议	我踩过的坑
2x	学习率×1.2	基本不用调
4x	学习率×1.5	偶尔需要warmup
8x	学习率×2.0	必须加warmup
16x	学习率×2.5	容易炸，建议用线性缩放

避坑指南：

我曾经在8卡机器上做分布式训练，每卡设了梯度累积步数=4，结果等效batch size变成了单卡的32倍。学习率没调，直接loss炸了。后来花了三天才排查出来。

1.4 实战代码与显存收益分析

直接上代码，这是我项目里常用的梯度累积实现：

import torch
import torch.nn as nn

class GradientAccumulator:
    def __init__(self, model, accumulation_steps=4):
        self.model = model
        self.accumulation_steps = accumulation_steps
        self.current_step = 0
        
    def backward(self, loss):
        # 缩放loss，保证梯度大小一致
        loss = loss / self.accumulation_steps
        loss.backward()
        
        self.current_step += 1
        
        if self.current_step % self.accumulation_steps == 0:
            # 累积够了，更新参数
            nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            optimizer.step()
            optimizer.zero_grad()
            
    def reset(self):
        self.current_step = 0
        optimizer.zero_grad()

这段代码里有个细节：loss除以了累积步数。为什么？因为如果不除，累积N步后的梯度就是正常梯度的N倍，参数更新步长会变大，模型容易震荡。

来看看显存收益分析。我用一个实际案例说明：

配置	显存占用	训练速度	模型效果
batch size=32（直接跑）	24GB（OOM）	-	-
batch size=8，累积4步	8GB	慢15%	基本一致
batch size=4，累积8步	4.5GB	慢30%	略有下降

看到没？显存从24GB降到了4.5GB，降了80%以上。代价是训练时间多了30%。这个trade-off，在显存紧张的时候非常值得。

我的个人经验：

一般建议累积步数不超过8步。超过8步后，梯度累积的误差会明显增大，模型效果开始下降。如果非要累积更多步，建议配合梯度裁剪和warmup一起用。

最后说一个容易被忽略的点：梯度累积会影响Batch Normalization的行为。因为BN是在每个微批次上计算的，累积步数多了，BN的统计量会不准。我一般会在累积步数大于4时，把BN换成Layer Normalization或者Group Normalization。

嗯，梯度累积这个技巧，说白了就是用小显存跑大batch。虽然有点trick的味道，但在实际工程中非常实用。我经手的项目里，至少有一半都用到了这个技术。

这张图很直观：左边是正常流程，右边是梯度累积流程。区别就在于多了一个"累积梯度"的循环。每次反向传播后不急着更新参数，而是把梯度存起来，等攒够了再更新。

想要更多资料，联系可可