A big batch gives a smoother, less noisy gradient and lets you train at the batch size a recipe was tuned for — but the activation memory of a forward/backward pass scales with the number of samples in it, so a batch of B can simply not fit on your GPU. Gradient accumulation fakes the big batch on small memory: split the effective batch into N small micro-batches, run forward and backward on each, and let the gradients pile up. I built a demo that proves the equivalence numerically. Here's the idea.
.grad already accumulates
The mechanism is free. PyTorch's .backward() adds into each parameter's .grad; it never overwrites. Normally you hide that by calling zero_grad() every step. Accumulation just… doesn't — it lets several backward passes pile into the same buffer.
loss_a.backward() # p.grad = g_a
loss_b.backward() # p.grad = g_a + g_b <-- ADDED, not replaced
# skip zero_grad() and gradients accumulate across micro-batches
Delay the step over N micro-batches
Feed N micro-batches of size m, backward on each, but only call optimizer.step() once every N. Between micro-batches the buffer grows; only one micro-batch's activations (m) is ever alive, so peak activation memory stays flat — independent of N.
for i, (xb, yb) in enumerate(loader): # micro-batches of size m
criterion(model(xb), yb).backward() # add into .grad
if (i + 1) % N == 0: # every N micro-batches...
opt.step(); opt.zero_grad() # ...one update, then reset
Scale by 1/N — the step everyone forgets
Each micro-batch loss is already a mean over m. Summing N of them gives N× the mean over m·N — too big by N. Divide each loss by N and the accumulated buffer becomes the true mean over the effective batch, matching a real big batch of B = m·N bit-for-bit, because the mean of equal-size means is the overall mean. Forget it and you've silently multiplied your learning rate by N.
loss = criterion(model(xb), yb) / N # <-- scale so the SUM is a MEAN
# without /N: buffer = Σ mean_j = N × true mean (the #1 accumulation bug)
# with /N: buffer = Σ mean_j / N = the big-batch mean ✓
What it costs, and what it doesn't
Total FLOPs are unchanged — you process the same B samples either way — so what you spend is time: N sequential smaller passes per update, and updates N× less often. That's the sharp contrast with Day-48 gradient checkpointing, which adds a recompute pass. Checkpointing trades compute for activation memory at fixed batch; accumulation trades steps/time for batch size at fixed memory — and the two compose cleanly.
The loop you'll actually write
Put it together and it's five lines around your normal loop — the same weights as a big batch of m·N, at the memory of m. It pairs cleanly with mixed precision (autocast + GradScaler) and with Day-48 checkpointing.
opt.zero_grad()
for i, (xb, yb) in enumerate(loader): # micro-batches, size m
loss = criterion(model(xb), yb) / N # scale so the sum is a mean
loss.backward() # accumulate into .grad
if (i + 1) % N == 0:
opt.step(); opt.zero_grad() # one real update every N
The one real gotcha: BatchNorm
The equivalence is exact only for per-sample losses. BatchNorm normalizes over the current micro-batch of m, not the effective m·N, so with BN present accumulation is not identical to a true big batch. Reach for GroupNorm, LayerNorm, or SyncBN. Mind the tail too — if the loader isn't a multiple of N, flush the leftover group after the loop or you silently drop those samples.
The takeaway I now carry: a big batch is a memory problem, not a math one. .grad already sums, so hand it N micro-batches, remember the 1/N, step once — and you've bought the big batch for the price of one small one.
Set m and N, and watch the memory bar stay flat while the accumulated gradient matches the big batch to machine zero:
https://dev48v.infy.uk/dl/day49-gradient-accumulation.html
Top comments (0)