Gradient Accumulation

A single accelerator might only fit a microbatch of a handful of sequences at a few thousand tokens each before running out of memory. A well-tuned pretraining run, meanwhile, often targets a global batch size in the millions of tokens, chosen with the critical batch size in mind. Gradient accumulation is the mechanism that closes that gap without needing a proportionally larger GPU.

The mechanism

Instead of computing a gradient and immediately stepping the optimiser, run forward and backward passes on several microbatches in sequence, summing their gradients into the same buffers without touching the parameters. Only after the target number of microbatches has been processed does the optimiser take a single step, using the accumulated gradient. The model sees one effective gradient per optimiser step, computed from the full target batch, exactly as if that batch had fit in memory all at once, just spread across several sequential passes instead of one.

What it does and does not buy you

Accumulation trades memory for time, not memory for nothing. Each microbatch still runs sequentially on the same device before the next one starts, so total FLOPs and total wall-clock time scale with the number of accumulation steps; a run using eight-way accumulation to hit a target batch size takes roughly eight times as long, per optimiser step, as one that could fit the whole batch at once. It does not reduce compute, it only makes a target batch size achievable on the hardware you actually have. Real pretraining setups combine it with data parallelism: replicas across GPUs handle part of the target batch in parallel, and each replica uses accumulation to cover whatever fraction still does not fit in a single microbatch, splitting the same target batch across both a parallel and a sequential axis.

Where it sits next to pipeline parallelism

Multi-stage pipeline parallelism already relies on streaming many small microbatches through the pipeline to keep every stage busy; Narayanan et al.'s Megatron-LM work describes the interleaved microbatch scheduling this requires. Gradient accumulation and pipeline microbatching are mechanically close cousins, both break one large intended batch into smaller sequential pieces, though they solve different problems: accumulation exists purely to fit a target batch in memory, pipeline scheduling exists to keep GPUs from sitting idle while waiting on the next pipeline stage.

The gotcha

The accumulated gradient must be divided by the number of microbatches (or, equivalently, the loss scaled down before each backward pass) before the optimiser step, to keep the effective learning rate consistent with what it would be at the target batch size. Forgetting this silently multiplies the effective learning rate by the accumulation step count, one of the more common and hard-to-notice bugs in a training pipeline, since the loss curve can still look plausible for a while before the mis-scaled updates cause visible instability.

When it falls down

• It does not fix a microbatch that is still too big. If even a single example does not fit in memory, that is an activation-memory problem, not a batch-size problem; it needs gradient checkpointing or activation offload, not accumulation.
• It buys a target batch size, not speed. Accumulation never makes training faster than running the same total tokens without it; it only makes an otherwise-impossible batch size reachable on limited hardware.
• Reduction consistency matters. Sum-then-divide and average-per-microbatch-then-average-again are not always numerically identical, and inconsistent handling across an accumulation split can make otherwise-identical runs diverge slightly.
• Poor overlap wastes the hardware. Deep accumulation without interleaving communication (gradient all-reduce, optimiser state sharding) with compute can leave GPUs idle between microbatches, eroding the very efficiency the technique is meant to preserve.

Further reading

• Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, arXiv:2104.04473 - microbatch scheduling in large-scale training pipelines.