Gradient Accumulation and Microbatches

The problem: the batch you want will not fit

Batch size is not a free knob. Larger batches give a less noisy estimate of the gradient, let the optimizer take more confident steps, and on modern hardware they improve arithmetic intensity so each matmul runs closer to peak throughput. Research recipes routinely quote effective batch sizes of hundreds of thousands or millions of tokens because that is where the training dynamics they tuned actually live. But the batch that helps optimization is often far larger than the batch that fits in memory.

The thing that fills memory during training is not the weights — it is the activations saved for the backward pass. Every layer stashes its inputs (and often intermediate tensors) so it can compute gradients on the way back, and that storage scales with the batch size and the sequence length. Double the batch and you roughly double the activation memory. So you hit a wall: a batch of 8 sequences fits, a batch of 256 does not, and yet 256 is the batch your recipe was tuned for. Gradient accumulation is the way out. It decouples the statistical batch size that the optimizer sees from the physical batch size that has to fit through the device at any one instant.

The core idea: gradients add up

The whole method rests on one fact from calculus: the gradient of a sum is the sum of the gradients. Training loss over a batch is (almost always) an average of per-example losses, and differentiation is linear, so the gradient of the batch loss is the average of the per-example gradients. Nothing forces you to compute that average in a single pass. You can compute the per-example gradients in chunks and add them together, and the result is identical.

Deep-learning frameworks make this natural because backward passes accumulate by default. When you call loss.backward(), each parameter’s .grad buffer is not overwritten — the freshly computed gradient is added to whatever was already there. Normally you hide this by calling zero_grad() before every backward, so each step starts from a clean slate. Gradient accumulation simply removes that reset for K-1 of every K microbatches: you let .grad keep growing across several backward passes, and only when the accumulation buffer holds the full batch’s gradient do you call optimizer.step() and then zero_grad(). The optimizer cannot tell whether that buffer was filled by one big batch or by K small ones.

The training loop, side by side

The change to your code is small — a loop counter, a conditional step, and a loss rescale — but every piece earns its place.

# Standard: one step per batch
for batch in loader:
    optimizer.zero_grad()
    loss = criterion(model(batch.x), batch.y)   # mean over the batch
    loss.backward()
    optimizer.step()

# Accumulated: one step per K microbatches
K = 8
optimizer.zero_grad()
for i, micro in enumerate(loader):
    loss = criterion(model(micro.x), micro.y)   # mean over the microbatch
    (loss / K).backward()        # scale so K of these average, not sum
    if (i + 1) % K == 0:
        optimizer.step()
        optimizer.zero_grad()    # reset only after the step

Three things carry the method. First, zero_grad() moves outside the inner cadence so gradients survive across microbatches. Second, the optimizer step fires only on every K-th microbatch, so K forward/backward passes feed one update. Third — the detail that trips everyone — the loss is divided by K before backward(). Without that division the accumulation buffer holds the sum of K mean-losses, which is K times too large, and your effective learning rate silently balloons by a factor of K.

The normalization detail: why divide by K

This is the single most common bug, so it is worth stating precisely. Suppose your loss for a microbatch is a mean over its examples — the default for most loss functions. Call the microbatch mean-losses L_1, L_2, ..., L_K. What you want the optimizer to see is the mean loss over the whole effective batch, which (when every microbatch has the same number of examples) is L = (1/K) Σ_k L_k. Its gradient is ∇L = (1/K) Σ_k ∇L_k.

If you call L_k.backward() K times without scaling, the accumulation buffer ends up holding Σ_k ∇L_k — the sum, which is K times the quantity you want. So you scale each microbatch loss by 1/K before its backward pass, and the buffer accumulates Σ_k (1/K)∇L_k = (1/K)Σ_k ∇L_k = ∇L, exactly. The mirror-image case matters too: if your loss uses reduction='sum' instead of mean, you do not divide by K — you divide by the total number of examples once, because the per-microbatch sums already add to the batch sum. Match the scale factor to the reduction, and the accumulated gradient equals the true large-batch gradient to the last bit.

A worked numeric example

Concreteness kills ambiguity. Take one scalar parameter w, an effective batch of 8 examples, and a microbatch size of 2, so K = 4. Say the eight per-example gradients ∂L_i/∂w happen to be:

g = [ 2, 4,   6, 8,   1, 3,   5, 7 ]

True full-batch gradient (mean over 8):
  g_full = (2+4+6+8+1+3+5+7) / 8 = 36 / 8 = 4.5

Now run it as four microbatches of two, each producing a mean loss, each scaled by 1/K = 1/4 before backward:

microbatch means:  (2+4)/2 = 3   (6+8)/2 = 7   (1+3)/2 = 2   (5+7)/2 = 6
scaled by 1/4:      3/4 = 0.75    7/4 = 1.75    2/4 = 0.50    6/4 = 1.50
accumulated sum:   0.75 + 1.75 + 0.50 + 1.50 = 4.5

The accumulated gradient is 4.5 — identical to the full-batch gradient, not approximately, but exactly. Drop the 1/4 scaling and the buffer would read 3+7+2+6 = 18, four times too big; the optimizer would take a step four times larger than intended, which is why an untuned accumulation setup often diverges or needs a mysteriously smaller learning rate. The arithmetic is the whole proof: accumulation is a regrouping of the same sum.

Effective batch size = microbatch x K x data-parallel

Once accumulation is in play, ‘batch size’ splits into several numbers that are easy to confuse. The one that governs training dynamics is the effective (or global) batch size, and on a distributed run it is a product of three factors:

B_eff = m x K x D

  m = per-device microbatch size   (what fits on one device at once)
  K = accumulation steps           (microbatches per optimizer step)
  D = data-parallel world size     (number of replicas / devices)

Each of the D replicas processes its own m x K examples and, at the optimizer step, the replicas average their gradients (an all-reduce), so the single update reflects all m × K × D examples. For language models people usually quote the effective batch in tokens: multiply by the sequence length S to get B_eff × S tokens per step. The practical consequence is a lever with three independent settings that all trade against each other. Want a bigger effective batch but out of memory? Raise K. Have more devices? Raising D gets you throughput and a bigger batch for free. Only m is bounded by per-device memory; K and D let you reach almost any global batch you like.

Accumulation and data parallelism together

Combining accumulation with data-parallel training exposes one efficiency subtlety worth understanding. In synchronous data parallelism (PyTorch DDP, for instance), every backward() normally triggers an all-reduce that averages gradients across replicas. That communication is expensive. If you naively accumulate over K microbatches, you would pay for K all-reduces per optimizer step when you only need one — the replicas only have to agree at the moment of the step.

The fix is to suppress synchronization on the intermediate microbatches and allow it only on the last one before the step. DDP exposes this as the no_sync() context manager: wrap the first K-1 microbatches in it so their gradients accumulate locally without communication, then run the K-th microbatch normally so a single all-reduce averages the fully accumulated local gradients. The result is the same mathematical update as one giant synchronized batch, at one-Kth of the communication cost. This is exactly why large-scale training loops lean on accumulation not only for memory but for communication amortization: fewer, larger syncs are cheaper than many small ones, and accumulation is the mechanism that batches them up.

Where BatchNorm breaks the equivalence

The exact equivalence has one important precondition: the model must be a per-example function, computing each example’s loss independently of the others in the batch. The moment a layer mixes information across the batch dimension, splitting the batch changes the computation, and accumulation stops matching a true large batch. BatchNorm is the canonical offender.

BatchNorm normalizes each activation using the mean and variance computed over the current batch. With a true batch of 256, those statistics are estimated from 256 examples. With accumulation over microbatches of 8, each forward pass computes mean and variance from only 8 examples — a noisier, differently distributed estimate — and there is no mechanism to pool statistics across microbatches, because they run in separate forward passes with no shared state. So accumulating K microbatches of 8 through BatchNorm is not equivalent to one batch of 8K; the normalization sees a much smaller effective sample, which shifts both the forward activations and the gradients. If you must accumulate with BatchNorm, you need a batch-independent variant — SyncBatchNorm (pool statistics across replicas), GroupNorm, or a running-statistics scheme — otherwise the ‘large batch’ you think you are simulating does not exist.

Why transformers are safe: LayerNorm and RMSNorm

Here is the good news for anyone doing this work in the transformer world: the BatchNorm caveat essentially does not apply. Transformers normalize with LayerNorm or RMSNorm, and both are computed per token, across the feature (hidden) dimension — never across the batch. LayerNorm subtracts each token’s own mean over its d features and divides by its own standard deviation; RMSNorm divides each token by its own root-mean-square over features and skips the mean entirely. Neither touches another example’s activations.

Because the normalization statistic of token i depends only on token i, a token’s forward pass and its gradient are identical whether it rode in a microbatch of 8 or a batch of 256. That is precisely the per-example property the equivalence needs, so for a standard transformer, accumulation over K microbatches is bit-for-bit the same update as the corresponding large batch (up to floating-point reduction order). This is why gradient accumulation is a first-class, everyday tool in LLM training and fine-tuning: the dominant architecture happens to use exactly the batch-independent normalization that makes the trick exact. Dropout and other per-token stochastic layers are fine too, as long as they do not couple examples.

Learning-rate scaling: the batch you simulate is real

Because accumulation genuinely enlarges the effective batch, it inherits the large-batch learning-rate question. A larger batch gives a lower-variance gradient estimate, which lets — and usually requires — a larger step. The two rules of thumb are the linear scaling rule (scale the learning rate in proportion to the batch size: double B_eff, double the LR), which works well for SGD with momentum, and the square-root rule (scale LR with √B_eff), often a better fit for Adam-family optimizers where the update is normalized by a running gradient magnitude.

The trap specific to accumulation is subtle: if you switch from an effective batch of 32 (no accumulation) to 256 (K=8) but leave the learning rate untouched, you are now taking small, timid steps relative to your batch, and training slows or stalls even though nothing looks broken. Conversely, forgetting the 1/K scale from earlier multiplies your effective LR by K on top of everything — a double error. Large batches also tend to need a longer warmup, because the confident early steps of a big batch can destabilize a freshly initialized model. Treat the accumulated batch as the real batch it is, and scale the LR (and warmup) accordingly.

The memory side: what you actually save

It helps to be precise about which memory accumulation reduces and which it does not. Training memory has four rough tenants: the parameters, the gradients (one buffer the size of the parameters), the optimizer state (for Adam, two more parameter-sized buffers — first and second moments), and the activations saved for backprop. The first three scale with model size and are completely independent of batch size. Only the activation memory scales with how many examples you push through at once.

Gradient accumulation shrinks exactly one of those tenants: activations. By running microbatches of size m instead of a batch of m × K, you hold only one microbatch’s worth of activations in flight at a time, cutting peak activation memory by a factor of K. The gradient buffer does not grow — gradients accumulate in place into the same parameter-sized buffer, so accumulating over 8 microbatches costs no more gradient memory than one. This is the elegant part: you get an 8× larger effective batch while your peak memory is set by the single microbatch plus the fixed model/optimizer overhead. For a memory-bound setup — a big model on a modest device, or a CPU box — that is often the difference between training and an out-of-memory crash.

The trade you are making: throughput for memory

Nothing is free. Accumulation buys memory headroom by spending wall-clock time. To reach one optimizer step you now run K forward passes and K backward passes instead of one of each. The total arithmetic is essentially unchanged — K microbatches of size m do the same FLOPs as one batch of mK — but you have serialized that work into K sequential launches rather than one wide, efficient pass.

That serialization has a real cost on parallel hardware. A larger physical batch keeps the compute units fuller (higher arithmetic intensity, better amortization of kernel-launch and memory-bandwidth overheads), so K small microbatches usually run somewhat slower in aggregate than one large batch would — if the large batch could fit, which is the whole reason you are here. The honest framing: accumulation is not a speedup, it is a feasibility tool. It lets you train at an effective batch size you could not otherwise reach, at the price of lower throughput than an imaginary device that had enough memory to do it in one pass. When you do have the memory, a bigger real microbatch beats accumulation; when you do not, accumulation is what makes the target batch possible at all.

Common pitfalls that quietly corrupt the run

Several failure modes look like normal training but silently break the equivalence. The missing 1/K scale is the headliner — already covered, and worth double-checking against your loss reduction. The variable-token problem is sneakier: in language modeling the loss is a mean over non-padded tokens, and microbatches rarely contain the same number of real tokens. Scaling each by a flat 1/K then weights a 400-token microbatch the same as a 900-token one, which is not the true token-mean of the whole batch. The rigorous fix is to sum the token-level losses, accumulate, and divide once by the total real token count across the K microbatches.

Other traps: forgetting to move zero_grad() outside the inner loop (so each microbatch wipes the last and you effectively train on microbatches, not the big batch); dividing the metric you log by K and misreading the loss curve; letting a stray BatchNorm or any cross-example op sneak in; and, with automatic mixed precision, stepping the gradient scaler on every microbatch instead of once per optimizer step. Each one produces a run that trains — just not the run you intended.

Why this matters most on small and CPU budgets

Gradient accumulation earns its keep precisely where memory is scarce, which makes it a staple of the small-model and CPU-training world this series cares about. On a single modest accelerator, or a CPU box with tens of gigabytes and no fast interconnect, the physical microbatch you can afford might be a handful of sequences — far below the effective batch a stable optimization recipe assumes. Accumulation closes that gap: set m to whatever fits, then choose K to hit the effective batch the recipe was tuned for.

It also composes cleanly with the other memory tricks. Stack it with activation checkpointing (recompute activations in the backward pass instead of storing them) to push m higher still, and with mixed precision to halve activation and buffer sizes. The mental model to keep: K and the data-parallel size D control the statistics of training (the effective batch and thus the optimization dynamics), while m, checkpointing, and precision control the memory footprint. Because transformers normalize per token, you can turn the memory knobs freely without disturbing the statistics — letting you reproduce a large-batch training recipe faithfully on hardware that could never hold that batch all at once. That decoupling is the quiet superpower of gradient accumulation.