4. 梯度累积：原理与实现，Batch Size与梯度累积的关系

说到梯度累积，我得先坦白一件事。

我刚开始做大模型训练那会儿，经常遇到一个让人抓狂的问题——明明代码写得没问题，显卡也够新，可一跑起来就报 OOM。后来我发现，很多时候不是显存不够，而是我们太贪心了，总想把 Batch Size 设得很大。

梯度累积，说白了就是「拆东墙补西墙」的聪明做法。它让你用小显存，跑出大 Batch Size 的效果。

4.1 为什么需要梯度累积？

先想一个问题：为什么大 Batch Size 训练效果好？

因为梯度更稳定，收敛更平滑。但大 Batch Size 意味着一次要加载更多数据，显存扛不住。

那怎么办？

梯度累积的思路很简单：

• 不一次性算完所有样本的梯度
• 而是分几次算，每次算一小批
• 把梯度累加起来，最后统一更新参数

举个例子。你想用 Batch Size = 64 训练，但显存只够跑 Batch Size = 16。那就分 4 步走：

1. 前 16 个样本，算梯度，存起来
2. 再 16 个样本，算梯度，累加上去
3. 再 16 个，再累加
4. 最后 16 个，累加完，统一更新参数

嗯，这就是梯度累积的核心逻辑。

4.2 Batch Size 与梯度累积的关系

这里有个容易踩的坑。我刚开始用梯度累积时，以为只要累积步数够多，效果就完全等价于大 Batch Size。其实不是。

为什么累积太多步会出问题？

因为梯度累积本质上是把多个小 Batch 的梯度「平均化」。如果累积步数太多，梯度更新频率变低，模型收敛会变慢。我遇到过有人把累积步数设到 32，结果训练了三天 loss 纹丝不动。

4.3 梯度累积的实现原理

实现起来其实不复杂。核心就三步：

• 前向传播：正常算 loss
• 反向传播：算梯度，但不更新参数
• 梯度累积：把梯度累加到缓存里
• 参数更新：累积够步数后，统一更新

这里有个关键点：每次反向传播后，梯度是累加的，不是覆盖的。PyTorch 默认会累积梯度，所以很多人不小心就踩了坑——忘了清零梯度。

4.4 实战代码示例

直接上代码。这是一个完整的梯度累积训练循环：

import torch
import torch.nn as nn
import torch.optim as optim

# 假设模型和数据
model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# 梯度累积参数
accumulation_steps = 4  # 累积 4 步
current_step = 0

for epoch in range(10):
    for batch_idx, (data, target) in enumerate(dataloader):
        # 1. 前向传播
        output = model(data)
        loss = criterion(output, target)
        
        # 2. 反向传播（梯度会累积）
        loss.backward()
        
        current_step += 1
        
        # 3. 累积够步数后，更新参数
        if current_step % accumulation_steps == 0:
            optimizer.step()      # 更新参数
            optimizer.zero_grad() # 清零梯度
            
            # 可选：打印 loss
            print(f'Epoch {epoch}, Step {current_step}, Loss: {loss.item():.4f}')

这段代码看着简单，但有几个细节要注意：

• loss 要除以累积步数：如果不除，累积后的 loss 会变大，梯度也会变大。我一般这样写：loss = loss / accumulation_steps
• 清零时机要对：一定要在 optimizer.step() 之后清零，否则梯度会一直累积下去
• 学习率要调整：累积步数越大，有效 Batch Size 越大，学习率可以适当调大一点

with torch.cuda.amp.autocast():
    output = model(data)
    loss = criterion(output, target)
    loss = loss / accumulation_steps

scaler.scale(loss).backward()

if (batch_idx + 1) % accumulation_steps == 0:
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

4.5 知识体系图

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

4.6 总结一下

梯度累积不是什么高深的技术，但它确实能解决很多实际问题。我个人的经验是：

• 累积步数 2-4 步最稳妥，效果和显存占用平衡得最好
• 记得把 loss 除以累积步数，不然梯度会偏大
• 混合精度训练时，注意 scaler 的用法
• 如果累积步数超过 8，建议调低学习率

说白了，梯度累积就是「用时间换空间」。虽然训练时间会变长一点，但至少不会 OOM 了。你想想看，与其让训练中断，不如多等一会儿，对吧？

想要更多资料，联系可可