DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

Gradient accumulation fakes a big batch on small memory — .grad already sums, so scale each micro-batch by 1/N and step once

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

Scale by 1/N — the step everyone forgets

Each micro-batch loss is already a mean over m. Summing N of them gives 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  ✓
Enter fullscreen mode Exit fullscreen mode

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 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
Enter fullscreen mode Exit fullscreen mode

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)