Gradient accumulation & the effective batch

The optimizer cares about one number — the effective batch size — and the hardware imposes another — the micro-batch that fits in HBM. Gradient accumulation is the lever that decouples them: run several micro-batches, sum their gradients, step once. It is also where DDP's communication can be quietly cut by a factor of G.

Two batch sizes that want to be different

There are two completely separate batch sizes in a training run, and conflating them is the most common source of "why won't my loss go down" confusion.

• The micro-batch b — how many samples one GPU runs through in a single forward/backward. This is a hardware number: it's bounded by HBM, because activations scale with b (lesson 10). Make it too big and you OOM.
• The effective (global) batch B_eff — how many samples contribute to one optimizer step. This is an optimization number: it sets the gradient's signal-to-noise ratio, and the right value is a property of the problem, not the GPU.

Data parallelism (lesson 04) already multiplies the micro-batch by the world size: N ranks each run b samples and AllReduce the gradients, so one step sees b·N samples. Gradient accumulation multiplies it again, in time:

B_eff = b · N · G

where G is the number of micro-batches accumulated before stepping. The point: if the optimizer wants B_eff = 4{,}096 but each GPU only fits b = 8 and you have N = 64 ranks, you're at 512 — short by 8×. Set G = 8 and you hit the target without buying more GPUs and without OOMing.

The mechanic — why summing gradients just works

A loss averaged over a batch is the mean of per-sample losses, and the gradient is linear in the loss, so the gradient of a big batch is the average of the gradients of its parts:

∇L(B) = (1/G) · Σ_{g=1}^{G} ∇L(micro-batch g)

So you don't need all G·b samples in memory at once. Run micro-batch 1, backward, add its gradients into the .grad buffers (PyTorch's backward() accumulates by default — this is why you call zero_grad()). Run micro-batch 2, add. … After G of them, the buffer holds the summed gradient; scale by 1/G, take one optimizer step, zero the buffers. The activations of micro-batch g are freed before g{+}1 starts, so peak activation memory is set by b, not G·b.

One precision detail: accumulate the gradients in fp32. The reason is bf16's mantissa: it carries only ~7 mantissa bits, so once a running sum grows past a micro-batch gradient by more than ~2⁷ in magnitude, that small addend rounds away entirely — sum G bf16 partials into a bf16 buffer and you silently drop the low-order contributions of the later (or smaller) micro-batches. A running fp32 accumulator has ~23 mantissa bits, enough headroom to keep those small addends, so the standard practice is to keep the .grad accumulator in fp32. (This is the same fp32 master-gradient buffer that mixed-precision training maintains for the optimizer step — detailed in lesson 17 — but the argument here stands on its own: it is purely about not losing bits when you sum many small numbers.)

Interactive · compose the effective batch

Three multipliers, one product. Slide the micro-batch (bounded by what fits in HBM), the data-parallel world size, and the accumulation depth. The grid shows every sample that contributes to a single optimizer step, coloured by which rank and which accumulation pass produced it. The memory bar tracks only b — notice that cranking G grows the effective batch without moving the memory bar at all.

The free win — skipping AllReduce with no_sync

Here is the systems subtlety. Naïve DDP fires an AllReduce on every backward() — it overlaps the reduction with the backward compute (lesson 04). But during accumulation, the gradients of micro-batches 1 … G{-}1 are intermediate: they're going to be summed locally anyway. AllReducing each one is pure waste — you'd average a partial sum across ranks, then keep adding to it, then average again. The mathematically identical and far cheaper plan is: accumulate locally and silently for the first G{-}1 passes, AllReduce only on the G-th.

PyTorch exposes this as the model.no_sync() context manager (and FSDP has an equivalent). The effect on communication is exactly G×: one AllReduce per optimizer step instead of one per micro-batch.

How big should the effective batch be? The critical batch size

Bigger batch → less gradient noise → you can take a larger, more confident step → fewer optimizer steps to reach a target loss. So why not make B_eff enormous? Because the returns saturate. McCandlish et al. (2018) formalised this: there is a critical batch size B_crit set by the gradient's noise-to-signal ratio, and

steps(B) ≈ S_min · (1 + B_crit / B) examples(B) ≈ E_min · (1 + B / B_crit)

Read the two curves against each other. Below B_crit, gradient noise dominates: doubling the batch nearly halves the number of steps — you're "buying" wall-clock speedup almost for free (each step processes more data but the step count drops proportionally). Above B_crit, the gradient is already clean: doubling the batch barely reduces steps but doubles the data (and compute) you burn per step. B_crit is the knee — the most parallelism you can exploit before you're just wasting FLOPs. Crucially B_crit is not static: it tracks the gradient noise scale, which rises as training proceeds (the loss falls, the gradient signal shrinks relative to its per-sample noise, so a larger batch is needed to average that noise down) — which is why large runs ramp the batch size over training rather than fixing it, starting small while the gradient is informative and growing it as the knee moves right.

Linear LR scaling — the batch and the learning rate move together

You cannot change B_eff and leave the learning rate alone. The widely-used rule of thumb (Goyal et al., 2017, "1 hour ImageNet"): in the small-batch regime, scale the learning rate linearly with the batch size — double B_eff, double the LR — because a larger batch gives a lower-variance gradient that a bigger step can safely exploit. Two caveats that matter in practice:

• Warmup is mandatory at large batch. A large LR applied to the noisy, ill-conditioned dynamics of the first few hundred steps blows up. Linearly ramp the LR from ~0 over a warmup window, then apply the scaled value.
• Linear scaling holds only below B_crit. Past the knee, the gradient is no longer noise-limited, the variance argument breaks, and people switch to square-root scaling or just stop increasing the batch. This is the same knee the previous widget draws.

Putting it together — the sizing recipe

1. Pick B_eff from the optimization side — at or just below B_crit for the model/data. This is a convergence decision, not a hardware one.
2. Find the largest micro-batch b that fits in HBM after FSDP + activation checkpointing (lessons 05, 20) — that's the memory budget talking.
3. Read off your DP world size N from the cluster you have.
4. Solve G = B_eff / (b·N), rounding to keep B_eff on target. Turn on no_sync so the G{-}1 intermediate AllReduces disappear.
5. Set LR by linear scaling from a known-good small-batch baseline, with warmup.