Gradient Accumulation
Decoupling batch size from GPU memory
Larger training batches generally produce more stable gradient estimates and are often needed to match published training recipes or to use higher learning rates safely. But batch size is directly constrained by available accelerator memory — every example in a batch needs its own activations stored for the backward pass. Gradient accumulation breaks this coupling: instead of computing gradients from one large batch in a single forward/backward pass, it splits the desired batch into several smaller micro-batches, computes and accumulates (sums) their gradients across multiple forward/backward passes, and only performs the actual weight update after accumulating gradients from all the micro-batches — mathematically equivalent to having used the full batch size, at a fraction of the peak memory.
Why it's mathematically equivalent
The gradient of a loss averaged over a batch is itself the average of the gradients computed on each example (or sub-batch) individually, by linearity of differentiation: $\nabla_\theta \frac{1}{N}\sum_i \mathcal{L}_i = \frac{1}{N}\sum_i \nabla_\theta \mathcal{L}_i$. So computing gradients on $k$ micro-batches of size $m$ and averaging them is identical to computing one gradient on a single batch of size $k \times m$ — the model has no way of distinguishing between the two, since the final accumulated gradient used for the weight update is the same either way. The only thing that changes is wall-clock time (more, smaller forward/backward passes instead of one large one) and peak memory (much lower, since only one micro-batch's activations are held at a time).
accumulation_steps = 8 # simulate a batch 8x larger than fits in memory
optimizer.zero_grad()
for step, micro_batch in enumerate(dataloader):
outputs = model(micro_batch["input"])
loss = loss_fn(outputs, micro_batch["target"])
# Scale down so summed gradients equal the average over the full batch
loss = loss / accumulation_steps
loss.backward() # gradients accumulate in .grad by default
if (step + 1) % accumulation_steps == 0:
optimizer.step() # apply the accumulated gradient, as if batch was 8x larger
optimizer.zero_grad() # reset for the next accumulation cycleBatch normalization breaks the equivalence
Large batch vs. gradient accumulation vs. small batch
| Setup | Peak memory | Gradient quality | Wall-clock time per effective batch |
|---|---|---|---|
| Large batch, single pass | High | Stable (full batch statistics) | Fastest, if it fits in memory |
| Gradient accumulation | Low — one micro-batch at a time | Equivalent to large batch (for most layer types) | Slower — multiple sequential passes |
| Small batch, no accumulation | Low | Noisier gradient estimates | Fast per step, but may need more steps to converge well |
Accumulation also helps in distributed settings
A simple but essential tool
Compared to the more involved infrastructure of pipeline, tensor, or ZeRO-style parallelism, gradient accumulation is almost embarrassingly simple — a few lines of training-loop logic — yet it's one of the most commonly reached-for tools for fitting a target effective batch size on constrained hardware, precisely because it requires no special infrastructure, multi-device setup, or framework beyond ordinary autograd.
What's next?
Gradient accumulation manages memory within a single training step; the next concern for any long-running training job is managing failure and recovery across potentially thousands of steps spanning days or weeks, which is where disciplined checkpointing strategies become essential.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.