梯度累加（Gradient Accumulation）的代价：如何在显存受限时模拟大 Batch

在深度学习的训练过程中，Batch Size（批次大小）是一个至关重要的超参数。通常情况下，更大的 Batch Size 能够提供更准确的梯度估计，有助于模型收敛到更优的解。然而，当模型参数量巨大或输入数据维度极高时，有限的显存（VRAM）往往成为使用理想 Batch Size 的瓶颈。

梯度累加（Gradient Accumulation）是一种优雅的解决方案，它允许我们在不增加单次迭代显存占用的前提下，模拟出超大 Batch Size 的训练效果。

梯度累加的核心原理

梯度累加的核心思想基于梯度运算的线性特性。对于标准的随机梯度下降（SGD），一次优化器步骤是基于一个完整批次（Batch）的平均梯度来更新权重的。

如果我们有一个理想的 Batch Size $B_{eff}$，但由于显存限制，我们只能使用物理 Batch Size $B_{phy}$，其中 $B_{eff} = N imes B_{phy}$。

梯度累加的做法是：

1. 将 $B_{eff}$ 拆分为 $N$ 个小的 $B_{phy}$ 批次。
2. 对每个小批次执行前向传播和反向传播，但不立即执行权重更新（即不调用 optimizer.step()）。
3. 反向传播产生的梯度会被累加到模型参数的 .grad 属性中。
4. 当累积了 $N$ 个小批次的梯度后，执行一次完整的 optimizer.step()。

这样，单次权重更新时使用的梯度，就是 $N$ 个小批次梯度的平均值（如果损失函数经过了适当的归一化），这在数学上等价于使用了 $B_{eff}$ 的超大 Batch Size。

PyTorch 实战：梯度累加实现

在 PyTorch 中实现梯度累加的关键点在于控制 optimizer.step() 的调用时机，并在计算损失时进行适当的归一化。

关键步骤

1. 定义累积步长 (Accumulation Steps): 决定我们要模拟扩大多少倍 Batch Size。
2. 损失归一化: 由于我们是在 $N$ 次迭代中累加梯度，如果不将每次计算的 loss 除以 $N$，则相当于学习率被隐式放大了 $N$ 倍。因此，推荐将 loss 除以 accumulation_steps 再执行 backward()。
3. 条件更新: 使用取模运算 (%) 判断是否达到累积步长，然后调用 optimizer.step() 和 optimizer.zero_grad()。

可运行代码示例

以下是一个使用 PyTorch 演示梯度累加的简化训练循环：

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

# --- 配置参数 ---
accumulation_steps = 4 # 模拟扩大 Batch Size 4倍
physical_batch_size = 8 # 受限于显存的物理 Batch Size

effective_batch_size = physical_batch_size * accumulation_steps
print(f"等效 Batch Size: {effective_batch_size}")

# --- 模拟环境 ---
model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# 模拟数据加载器，总共包含20个小 Batch
dummy_data_loader = [(torch.randn(physical_batch_size, 10), 
                      torch.randn(physical_batch_size, 1)) 
                     for _ in range(20)]

# --- 梯度累加训练循环 ---
model.train()

for i, (inputs, targets) in enumerate(dummy_data_loader):
    # 1. 前向传播
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # 2. 损失归一化：确保累积梯度的平均性
    loss = loss / accumulation_steps

    # 3. 反向传播：梯度被累加到 .grad 属性中
    loss.backward()

    # 4. 条件更新：判断是否达到累积步长
    if (i + 1) % accumulation_steps == 0:
        print(f"[Step {i + 1}] 达到累积步长，执行权重更新和梯度清零")
        optimizer.step()
        optimizer.zero_grad()

# 5. 处理剩余梯度：如果数据集大小不是 accumulation_steps 的整数倍
# 确保在 Epoch 结束时，所有累积的梯度都被使用
if (len(dummy_data_loader) % accumulation_steps) != 0:
    print("处理最后的残余梯度")
    optimizer.step()
    optimizer.zero_grad()

成本与代价

虽然梯度累加成功解决了显存限制的问题，但它并非没有成本：

1. 训练时间增加： 虽然权重更新的次数减少了，但需要重复执行 $N$ 次前向和反向传播。这意味着 I/O 和计算开销的时间片会重复 $N$ 次，总训练时间会相应增加。
2. Batch Norm 影响： 如果模型中使用了 Batch Normalization (BN) 层，BN 层计算的均值和方差仍然是基于物理 Batch Size $B_{phy}$，而不是等效 Batch Size $B_{eff}$。这可能会导致 BN 统计数据不够稳定，从而影响模型的精度。解决这一问题的方法包括使用 Group Normalization 或同步 Batch Normalization（SyncBN）。