第三节：梯度累积与显存优化

梯度累积，说白了就是「攒一波再更新」。

我刚开始接触大模型训练时，最头疼的就是显存不够用。明明想跑个 32 的 batch size，结果 8 就爆显存了。后来发现，梯度累积这个技巧，简直就是显存不够时的救命稻草。

梯度累积的原理

先说说常规训练是怎么做的。你喂一个 batch 的数据，前向传播算 loss，反向传播算梯度，然后优化器立刻更新参数。这个过程，每个 batch 都要存一份完整的中间结果。

梯度累积的思路完全不同：

先跑几个小 batch，只算梯度，不更新参数
把梯度累加起来
攒够了，一次性更新参数

这样做的好处很明显——显存占用只取决于单个小 batch 的大小，而不是累积后的总 batch size。

核心公式：

有效 batch size = 单步 batch size × accumulation_steps

显存占用 ≈ 单步 batch size 的显存

举个例子。你想跑有效 batch size = 64，但单步只能塞下 8。那设置 accumulation_steps = 8 就行了。显存占用还是 8 的水平，但效果相当于 64。

accumulation_steps 怎么设？

这个参数其实很好理解。它告诉程序：「你攒够多少步，再更新一次参数」。

我个人习惯这样算：

# 假设目标有效 batch size = 64
# 单步最大 batch size = 8
accumulation_steps = 64 // 8  # 结果是 8

嗯，这里要注意一点。accumulation_steps 必须是整数，而且最好能被有效 batch size 整除。不然最后一个累积步数会少，影响训练稳定性。

我的经验：

accumulation_steps 一般设 2、4、8 比较常见。超过 16 的话，训练速度会明显变慢，因为参数更新频率太低了。

与 Batch Size 的关系

很多人搞混一个概念：梯度累积 ≠ 增大 batch size。

为什么？

你想想看，真正的 batch size 增大，是所有样本的梯度同时算出来的。而梯度累积是分步算的，每一步的梯度基于不同的数据子集。虽然数学上等价，但实际训练中，由于 BN 层的统计量不同，效果会有细微差别。

我曾经踩过一个坑：把 accumulation_steps 设得特别大，比如 64，结果模型收敛变慢了。后来发现是 BN 层的 running mean 和 running var 更新频率太低，导致统计量不准。

配置	显存占用	训练速度	模型效果
batch_size=64, 无累积	高（可能爆显存）	快	好
batch_size=8, accum=8	低	中等	接近
batch_size=4, accum=16	很低	慢	略差

实战代码示例

下面是我在实际项目中用过的梯度累积代码。以 PyTorch 为例：

# 梯度累积训练循环
accumulation_steps = 8
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 关键：除以 accumulation_steps，保持 loss 尺度一致
    loss = loss / accumulation_steps
    loss.backward()
    
    # 每 accumulation_steps 步更新一次
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 别忘了处理最后不足一个累积步的情况
if (i + 1) % accumulation_steps != 0:
    optimizer.step()
    optimizer.zero_grad()

避坑指南：

我曾经忘记把 loss 除以 accumulation_steps，结果 loss 直接炸了。因为累积了 8 步的梯度，更新时的梯度大小是正常的 8 倍，参数一步就飞出去了。

另外，如果你用分布式训练，梯度累积要配合梯度同步。我一般这样处理：

# 分布式 + 梯度累积
from torch.nn.parallel import DistributedDataParallel as DDP

model = DDP(model)
accumulation_steps = 8

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:
        # 只在更新时同步梯度
        for param in model.parameters():
            if param.grad is not None:
                # 手动触发梯度同步
                pass
        optimizer.step()
        optimizer.zero_grad()

知识体系图

下面这张图展示了梯度累积的核心逻辑：

总结一下：

梯度累积不是银弹。它解决的是显存瓶颈，不是计算瓶颈。如果你的显存够用，直接跑大 batch 更快。但如果你像我一样，经常在 24G 显存的卡上跑大模型，那梯度累积就是必备技能。

最后提醒一句：别忘了调学习率。有效 batch size 变了，学习率也要跟着调。我一般按比例缩放：batch size 翻倍，学习率也翻倍。但这个不是绝对的，具体还得看你的任务。