Gradient accumulation 怎么做？
为什么 PyTorch 默认累加 gradient？

英文问法Why does PyTorch accumulate gradients by default? · Implement gradient accumulation. · Is it equivalent to training with a larger batch?

考过面经：NVIDIA 框架岗 2023 onsite 第 3 轮（实现 backprop 后追问实现 grad-accumulation，知乎）· NVIDIA DL 全职 VO 2022（讨论"为什么by default grad 要累计"）

先弄清楚：每个概念是什么

.grad 累加: 每次 backward() 都把新 gradient 加到 parameter 的 .grad 上，不覆盖。例：同一个 batch 连续 backward 两次，.grad 是一次的 2 倍。
zero_grad(): 每次 optimizer step 之后把 .grad 清零，否则上一步的 gradient 会混进这一步。
Micro-batch 与 gradient accumulation: GPU memory 只放得下 \(b\) 个样本时，把大 batch 拆成 \(m\) 个 micro-batch，逐个 forward + backward，攒满 \(m\) 个才 step 一次。activation memory 只按 \(b\) 算。例：\(b = 8\)，\(m = 4\)，DP 卡数 \(N = 2\) → global batch \(B = 64\)。
Loss ÷ m: 每个 micro-batch 的 loss 是 \(b\) 个样本的 mean，\(m\) 个 gradient 直接相加是大 batch 的 \(m\) 倍，所以每个 loss 先除以 \(m\)。
no_sync(): DDP 每次 backward 都会 all-reduce gradient。前 \(m-1\) 个 micro-batch 放进 model.no_sync() 只在本卡累加，通信次数从 \(m\) 降到 1。FSDP 的 no_sync() 累加的是完整 gradient 而不是 shard，memory 更高。

看图：4 个 micro-batch 怎么攒成一次更新

图 1：b = 8、m = 4、N = 2，时间从上往下。右边小格是 .grad 里已经攒进去的份数。前 3 次 backward 在 no_sync() 里，第 4 次才 all-reduce（暖色框），然后 step、zero_grad()。

30 秒回答

backward() adds into .grad instead of overwriting it. Gradients are linear, so two backward calls give the gradient of the summed loss; that is what makes micro-batch accumulation and several separate losses work, at the price of calling zero_grad() every step. To accumulate, I split the batch into m micro-batches, divide each loss by m, call backward on each, then step once and zero the grads. With equal-sized micro-batches that gives exactly the big-batch gradient, with activation memory for only one micro-batch. It breaks for BatchNorm, whose statistics only see the micro-batch. Under DDP I wrap the first m − 1 backwards in no_sync(), so the all-reduce runs once.

推导 / 关键公式

\(L_k\) 是第 \(k\) 个 micro-batch（\(b\) 个样本）的 mean loss：

\[ L_B = \frac1B\sum_{i=1}^{B}\ell_i = \frac1m\sum_{k=1}^{m} L_k \] \[ \nabla L_B = \sum_{k=1}^{m}\nabla\Big(\frac{L_k}{m}\Big) \]

右边就是 \(m\) 次 backward 用 += 攒出来的。不除以 \(m\)，SGD 下等于 learning rate 放大 \(m\) 倍。加上 \(N\) 张卡的 DP，global batch：

\[ B = b\cdot m\cdot N \]

手算例子

\(y = wx\)，\(w = 0\)，\(x = (1, 2, 3, 4)\)，\(y = 2x\)，loss 取 \((wx - y)^2\) 的 mean，gradient \(= \text{mean}\big(2(wx - y)x\big)\)：

大 batch：\(-\tfrac{2}{4}(2 + 8 + 18 + 32) = -30\)
\(m = 2\)：前两个点得 \(-10\)，后两个点得 \(-50\)，直接累加 \(-60\)，多了一倍
每个 loss 先 ÷ 2：\(-5 + (-25) = -30\)

手撕代码（PyTorch + DDP）

from contextlib import nullcontext

def train_step(model, opt, micro_batches, loss_fn):
    m = len(micro_batches)
    for k, (x, y) in enumerate(micro_batches):
        last = k == m - 1
        # DDP：前 m−1 次 no_sync，只在本卡累加
        with nullcontext() if last else model.no_sync():
            loss = loss_fn(model(x), y) / m  # ÷ m
            loss.backward()                  # .grad +=
    opt.step()          # m 次 backward 后 step 一次
    opt.zero_grad()     # 清空，留给下一个 step

必答点

backward() 是 .grad +=：gradient 对 loss 是 linear 的，累加 = 对 loss 之和求导；代价是每步 zero_grad()
做法：loss ÷ \(m\) → backward \(m\) 次 → step() → zero_grad()
等价于 \(B = b\cdot m\cdot N\) 的大 batch，activation memory 只按 \(b\)
不等价：BatchNorm；micro-batch 大小（或 token 数）不同
DDP：前 \(m-1\) 次 no_sync()，all-reduce 从 \(m\) 次降到 1 次

追问

micro-batch 大小（或 token 数）不同怎么办？

÷ \(m\) 默认每个 micro-batch 权重一样，上例按 3 + 1 拆会得 \(-41.3\) 而不是 \(-30\)。应该用 sum 的 loss 除以整个 step 的样本数（LLM 用非 padding token 数）。Hugging Face 2024 年修过 Trainer 的这个 bug。

和 BatchNorm 一起用有什么问题？

BN 的 mean / variance 只来自 \(b\) 个样本，running statistics 也按 micro-batch 更新，可以换 GroupNorm。Transformer 用 LayerNorm / RMSNorm，不受影响（见 BatchNorm vs LayerNorm）。

加了 accumulation，learning rate 要改吗？

global batch \(B\) 不变就不用改，\(B\) 变大了就按新的 \(B\) 重调。lr schedule 和 warmup 按 optimizer step 计数。

zero_grad() 设成 None 和设成 0 有什么区别？

PyTorch 2.x 默认 set_to_none=True，省 memory。gradient 为 None 的 parameter optimizer 直接跳过，为 0 时仍然会 step（momentum 还会推着走）。

来源torch.Tensor.backward · PyTorch recipe: Zeroing out gradients · Optimizer.zero_grad（set_to_none）· DistributedDataParallel（no_sync）· FSDP（no_sync 的 memory note）· Narayanan et al. 2021 §3.3.1（p.5，\(m = B/(d\cdot b)\)）· Hugging Face, Fixing Gradient Accumulation（2024-10-16）

Gradient accumulation 怎么做？为什么 PyTorch 默认累加 gradient？