Learn AI Create the Future — with Pranay
Main site

Autograd in Depth

Gradient accumulation for large batches

Because gradients add up across backward calls, several small batches can impersonate one big batch — the standard trick when the batch you want will not fit in memory.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Gradient accumulation means processing a big batch in small instalments, adding up the blame, and updating the model once at the end.

Think of carrying twenty bricks upstairs when you can lift five. You make four trips, stacking bricks at the top each time, and only after the last trip do you start building. Nobody builds a quarter-wall after each trip. The stack at the top is the accumulated gradient; the building step is the weight update.

Why it exists

Bigger batches give steadier learning signals, and some training recipes assume batches far larger than a small GPU can hold. Memory runs out; the recipe does not care.

Here PyTorch's own behaviour becomes the tool. Gradients from every backward call add into .grad — they never overwrite. Usually you clear between steps. Skip the clearing on purpose, and four small batches deposit their blame into the same pile. Update once from the pile, and the maths matches one big batch.

How it works

big batch (does not fit):    [ 8 samples ] -> blame -> update

instalments (fit easily):    [2] -> blame ┐
                             [2] -> blame ├ pile up in .grad
                             [2] -> blame │
                             [2] -> blame ┘ -> ONE update, then clear

The model sees identical mathematics either way. Only your memory bill changes — and the time: four trips take longer than one.

A real example you have seen

Many well-known open-source models were trained with exactly this trick — research groups without giant clusters matching the batch sizes of rich labs, instalment by instalment. When a paper says "effective batch size 1024 on a single GPU", accumulation is the quiet machinery behind the sentence.

Remember this

  • Gradients add up across backward calls; clearing between steps is a choice.
  • Skip the clearing, update once per pile: small batches impersonate big ones.
  • Same maths, less memory, more time — a trade, not a free lunch.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Outputs verified with torch 2.5.1, CPU.

Proving the equivalence, not assuming it

accumulate.py
import torch

torch.manual_seed(0)
data = torch.randn(8, 4)
target = torch.randn(8, 1)

def fresh_model():
    torch.manual_seed(1)              # identical starting weights each time
    return torch.nn.Linear(4, 1)

# One big batch of 8
big = fresh_model()
loss = torch.nn.functional.mse_loss(big(data), target)
loss.backward()

# Four micro-batches of 2, gradients left to pile up
small = fresh_model()
for i in range(0, 8, 2):
    part = torch.nn.functional.mse_loss(small(data[i:i+2]), target[i:i+2])
    (part / 4).backward()             # divide by the number of micro-batches

print(torch.allclose(big.weight.grad, small.weight.grad, atol=1e-7))
print(big.weight.grad)
print(small.weight.grad)
Output
True
tensor([[-0.8099, -0.0660,  0.0155,  0.3777]])
tensor([[-0.8099, -0.0660,  0.0155,  0.3777]])

The walkthrough

The / 4 is the entire subtlety. Each micro-loss is a mean over its 2 samples. Four means, added, equal four times the big batch's mean — unless you divide. Divide each instalment by the number of instalments and the pile matches exactly. Forgetting this division is the classic accumulation bug: gradients 4× too large, which behaves like a 4× learning rate nobody asked for.

Identical to the seventh decimal. Not "roughly the same" — allclose at atol=1e-7 passes because the mathematics is the same sum, reordered. (Reordered float sums can differ in the last bits; that is the reason for the tolerance at all.)

In a real training loop, the pattern adds two lines to what you know:

training_pattern.py
accum_steps = 4
for step, (inputs, labels) in enumerate(batches):
    loss = loss_fn(model(inputs), labels) / accum_steps
    loss.backward()                          # deposits blame, no update yet
    if (step + 1) % accum_steps == 0:
        optimizer.step()                     # build the wall
        optimizer.zero_grad()                # clear the pile

(Sketch, not runnable alone — it needs your model and batches. The two moved lines are the point: step and zero_grad live inside the if.)

Common mistakes

Scaling nothing, or scaling twice. Divide the loss by accum_steps once. Dividing and also averaging gradients later, or dividing when your loss already uses reduction="sum", silently rescales learning. Decide where the division lives and write a comment saying so.

Uneven final group. If the data ends mid-pile, the leftover instalments got divided by a wrong count. Either drop the tail batch, or divide by the actual number of instalments in this pile.

Assuming everything scales like the gradients. Batch-norm statistics are computed per micro-batch — four instalments of 2 see noisier statistics than one batch of 8. Accumulation matches the gradients, not the batch-norm behaviour. Models heavy on batch-norm genuinely differ between the two setups; transformers (layer-norm) do not care.

Metric logging inside the pile. Logging loss per micro-step shows the divided values — a mysterious 4× drop in your curves. Multiply back, or log per update.

Piling for hundreds of steps. The gradient pile is fine, but stale piles interact badly with learning-rate schedules and gradient clipping applied per micro-step. Clip once per update, on the completed pile — see gradient clipping.

Try it yourself

Break the demo on purpose: remove the / 4 and print both gradients — confirm the ratio is exactly 4. Then set accum_steps so micro-batches of 2 impersonate a batch of 16 and verify against a real batch of 16, using the same fresh_model seeding trick.

What to learn next

Researcher — Mathematics and papers.

Why it is exact

For a mean-reduced loss over batch B split into K equal micro-batches B_k:

∇L_B = ∇[(1/K) Σ_k L_{B_k}] = (1/K) Σ_k ∇L_{B_k}

Where L_{B_k} is the mean loss over micro-batch k. Gradient linearity gives the identity; the accumulate-then-step procedure computes the right side term by term into .grad (the AccumulateGrad += of the leaf lesson). With unequal micro-batch sizes n_k, the correct weighting is Σ_k (n_k/n) ∇L_{B_k} — the "divide by K" shortcut assumes equal sizes.

The equivalence is exact for the gradient computation and fails only where the forward pass itself depends on batch composition: batch-norm's per-batch moments (Ioffe and Szegedy, 2015) — the practical fix being ghost/synchronised BN or norm layers without batch coupling — and any loss with cross-sample terms (contrastive losses over in-batch negatives lose negatives under splitting; InfoNCE quality degrades with micro-batch size).

Interaction with the optimiser, precisely

Accumulation changes nothing the optimiser sees: one .grad, one step. Therefore effective-batch-size heuristics apply to the accumulated batch — linear LR scaling with warmup (Goyal et al., 2017, Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour), and the square-root scaling regime for Adam-family optimisers under some conditions. Scheduler ticks, weight decay application, EMA updates and clipping all belong at update granularity, not micro-step granularity — auditing a codebase for "per-step things that should be per-update" is the review checklist for accumulation bugs.

Memory: activation memory drops from O(n) to O(n/K) while parameter, gradient and optimiser-state memory are unchanged — accumulation attacks exactly one of the four memory terms, which is why it composes with, rather than replaces, checkpointing (activations), mixed precision (all terms), and ZeRO-style sharding (states) (Rajbhandari et al., 2020, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models).

The statistics of batch size

Larger effective batches reduce gradient variance: Var[∇L_B] ∝ 1/n for i.i.d. samples. The gradient-noise-scale analysis (McCandlish et al., 2018, An Empirical Model of Large-Batch Training) predicts the critical batch size beyond which added samples stop buying wall-clock progress — accumulation lets a small device reach that regime, but also inherits its diminishing returns: past the critical size, K more trips buy noise reduction the optimisation no longer exploits. Generalisation effects of very large batches (Keskar et al., 2017, sharp minima; Smith et al., 2018 on LR/batch equivalences) apply to accumulated batches identically — the optimiser cannot tell the difference, which is the point.

References

  • Goyal et al. (2017), arXiv:1706.02677 — LR scaling, warmup, and the reference large-batch recipe.
  • McCandlish et al. (2018), arXiv:1812.06162 — gradient noise scale and critical batch size.
  • Ioffe and Szegedy (2015), Batch Normalization — the layer that breaks the equivalence.
  • Rajbhandari et al. (2020), arXiv:1910.02054 — the memory-term taxonomy accumulation slots into.

What to learn next