Gradient accumulation: effective batches without losing the correct normalization

What it is

Gradient accumulation processes several microbatches before updating model parameters. Their gradients contribute to one optimizer step, allowing a larger effective batch when a physical batch does not fit in device memory. PyTorch backward calls add to existing parameter gradients until those gradients are cleared. [1][2]

How it works

For equally sized microbatches and a mean loss over equally weighted examples, divide each loss by the number of microbatches in the accumulation window. Summing those gradients then yields the gradient of the combined mean loss. Unequal sizes require weighting by the actual contributing example count. For token losses, count valid, unmasked tokens; class weights or other reductions may require a different denominator. Averaging microbatch means blindly changes the objective. These normalization rules follow from differentiating the intended weighted mean.

Practical workflow

Clear gradients at the start of a window, run forward and backward for each normalized microbatch, then perform one optimizer step. Clear gradients for the next window. Separate forward passes normally do not require retaining earlier computation graphs. [2] Handle the final, incomplete window explicitly. With mixed precision, keep the gradient scale constant throughout accumulation, unscale once after the last backward pass, and then clip the combined gradient before stepping. Update the scaler at that boundary. [1] Advance a scheduler configured per optimizer update only when an update occurs.

Illustrative example

Suppose one device processes four microbatches of eight examples. The effective batch contains 32 examples, so each microbatch mean receives weight 1/4. A final window containing 8, 8, 8 and 2 examples contains 26 examples: the corresponding weights are 8/26, 8/26, 8/26 and 2/26. This is arithmetic, not a measured training result.

Limits and verification

Batch normalization uses statistics from each physical microbatch, so accumulation does not reproduce a single large batch automatically. [3] Random operations and floating-point summation order can also change results. Clipping each microbatch separately generally changes the combined update. The method reduces simultaneous activation storage; it does not eliminate model weights or optimizer states, and additional passes can increase runtime.

Before a long run, compare one accumulated update with one combined-batch update on a small deterministic model without batch-dependent layers. Inspect loss normalization, gradient differences within a suitable tolerance, update counts and partial-window behavior. These are recommended checks, not tests performed for this article.

Sources

[1] PyTorch: Automatic Mixed Precision examples [2] PyTorch: torch.Tensor.backward [3] PyTorch: BatchNorm1d