4、梯度累积策略：梯度累积原理、accumulation_steps设置、与Batch Size的关系、分布式下的梯度累积

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

我刚开始接触这个技巧时，觉得它挺反直觉的——明明可以一次算完，为什么要分多次？后来在训练一个超大模型时，显存直接爆了，才明白这玩意儿有多重要。

梯度累积原理

常规训练中，每个batch都会计算梯度，然后立即更新参数。梯度累积则不同：

• 前向传播 + 反向传播，计算梯度
• 不更新参数，把梯度存起来
• 重复N次，累积梯度
• 第N次时，用累积的梯度统一更新参数

嗯，这里要注意：累积的是梯度，不是损失。损失函数的值每次都会重新计算，但梯度会累加。

我在项目中遇到过一个问题：累积梯度后，学习率要不要调整？答案是——要。因为累积N步后，梯度总和变大了，如果学习率不变，参数更新步长会变大。我个人习惯把学习率乘以N，或者用Adam这种自适应优化器，效果更稳。

accumulation_steps设置

这个参数决定了「攒几次梯度再更新」。怎么设？我一般看显存说话。

举个例子：你本来想用batch size=64，但显存只够跑batch size=8。那accumulation_steps设成8，效果就相当于batch size=64了。

与Batch Size的关系

很多人以为「accumulation_steps × micro_batch_size = 等效batch size」。这个公式没错，但有个坑——Batch Normalization。

BN层在训练时，统计的是每个micro batch的均值和方差。如果你用accumulation_steps=8，micro_batch_size=8，等效batch size=64，但BN统计的还是8个样本的分布。这会导致：

• BN统计量不稳定
• 模型收敛变慢
• 推理时效果打折扣

怎么解决？我一般用以下方案：

• 用LayerNorm代替BN（Transformer模型天然适用）
• 或者用SyncBN（分布式下同步BN统计量）
• 实在不行，就增大micro_batch_size，减少accumulation_steps

分布式下的梯度累积

分布式训练 + 梯度累积，这里面的门道更多。我画了一张图帮你理解：

分布式下梯度累积，有两种主流策略：

策略一：每步同步

每个micro batch都做AllReduce。优点是梯度实时同步，缺点是通信开销大。我试过8卡训练，每步同步，通信时间占了总时间的40%。

策略二：累积后同步

每个GPU先本地累积N步，然后做一次AllReduce。通信次数减少到1/N，但需要更大的本地显存来存累积梯度。

# 分布式梯度累积
accumulation_steps = 8
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 累积梯度
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        # AllReduce 同步梯度
        for param in model.parameters():
            dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
            param.grad /= world_size
        
        optimizer.step()
        optimizer.zero_grad()

最后说一句：梯度累积不是银弹。它解决的是显存问题，但会引入额外的通信开销和数值精度问题。我一般只在显存实在不够用时才用，能跑大batch size就尽量跑大的。

想要更多资料，联系可可