Gradient Accumulation › Accumulation Correctness

Gradient accumulation introduces several subtle correctness issues that can silently undermine training if not addressed carefully. The mathematical equivalence to full-batch training holds only under specific conditions.

Batch normalization computes statistics (mean and variance) over the batch dimension. When you use gradient accumulation, each micro-batch has different statistics from the full logical batch. This means that:

1. The forward pass statistics used during training will differ from what they would be on the full batch.
2. The running mean and variance estimates updated during the forward pass will reflect micro-batch statistics rather than full-batch statistics.

This difference goes beyond numerical precision. Batch normalization layers compute a different function when the batch size changes, and the model's behavior during inference (which uses the running statistics) will reflect whichever batch size was used during training.

The magnitude of this effect depends on how small the micro-batches are. With micro-batches of 64, the batch statistics are already a reasonable estimate of the population statistics, and the deviation from full-batch behavior is small. With micro-batches of 2 or 4, the statistics are highly variable, and the normalization can be erratic.

Layer normalization, which normalizes over the feature dimension rather than the batch dimension, is not affected by this problem. This is one reason why transformer architectures, which use layer normalization, are more amenable to gradient accumulation than convolutional architectures that rely on batch normalization.

If you must use batch normalization with gradient accumulation, you can use ghost batch normalization: compute batch norm statistics as if the full logical batch were present by accumulating statistics across micro-batches before the normalization step. This is more complex to implement but preserves the intended behavior.

Dropout applies a random binary mask to activations during the forward pass, and the masked elements are excluded from the backward pass. When accumulating gradients across $K$ micro-batches, each micro-batch sees a different dropout mask. This is the correct behavior: the full logical batch processed in one shot would also apply dropout independently to each sample. The stochasticity is per-sample, not per-batch, so accumulation preserves the intended regularization.

However, a subtle issue arises if you set the random seed deterministically per micro-batch index rather than per sample. Ensure that your random state is not reset between micro-batches within an accumulation cycle.

As discussed in the gradient clipping material from the Neural Network Foundations part, gradient clipping constrains the gradient norm before the optimizer step to prevent exploding gradients. When using gradient accumulation, you must clip the gradient after all micro-batches have been processed, not after each one. Clipping after each micro-batch would artificially reduce the effective gradient magnitude and could prevent gradients from flowing correctly.

To understand why, consider a scenario where the true gradient norm is 2.0 and you are using a clip threshold of 1.0. With $K = 4$ micro-batches, each micro-batch contributes a gradient of norm approximately 0.5 (the full gradient divided by $K$). If you clip each micro-batch gradient at 1.0, no clipping occurs, but the accumulated gradient has norm 2.0, which would have been clipped. The accumulated gradient is twice the intended magnitude.

Conversely, if the true gradient norm is 0.25 and each micro-batch gradient has norm approximately 0.0625, clipping at 1.0 per micro-batch has no effect, but the result after accumulation (norm 0.25) is also fine. The correctness problem only manifests when the true full-batch gradient norm exceeds the clip threshold.

The correct procedure involves three steps performed in sequence:

1. Accumulate gradients across all $K$ micro-batches.
2. After the final backward pass, call gradient clipping on the accumulated gradients.
3. Then call the optimizer step.

Mixed precision training uses float16 (or bfloat16) for forward and backward passes to reduce memory usage and increase throughput. Float16 has limited dynamic range, and gradients that are very small can underflow to zero. To prevent this, a loss scaler multiplies the loss by a large factor $S$ before the backward pass and then divides the resulting gradients by $S$ before the optimizer step.

When combining gradient accumulation with mixed precision, the loss scale $S$ and the accumulation scaling factor $1/K$ interact. The combined scaling applied to each micro-batch loss before backpropagation is:

$$\mathcal{L}_k^{\text{scaled}} = \frac{S}{K} \cdot \mathcal{L}(\theta; B_k)$$

where:

• $\mathcal{L}_k^{\text{scaled}}$: the scaled loss passed to .backward() for micro-batch $k$
• $S$: the loss scale factor used by the mixed precision scaler (typically a large power of 2)
• $K$: the number of accumulation steps
• $\mathcal{L}(\theta; B_k)$: the unscaled mean loss over micro-batch $B_k$

After all $K$ micro-batches, the gradients are divided by $S$ before the optimizer step to reverse the loss scaling. Most modern training frameworks handle this automatically when you use their built-in gradient accumulation support. If implementing manually, ensure that both scaling factors are applied correctly and that the loss scaler's overflow detection operates on the accumulated gradients rather than on individual micro-batch gradients.

Bfloat16 (the 16-bit brain float format used in TPUs and supported by A100 GPUs) has wider dynamic range than float16 and often eliminates the need for loss scaling entirely. When using bfloat16, you can simplify the combined scaling: only the $1/K$ accumulation factor is required.

In data-parallel distributed training, each GPU processes a different mini-batch and gradients are averaged across GPUs before the optimizer step. This all-reduce communication typically happens automatically during the backward pass via hooks attached to each parameter.

When using gradient accumulation in a distributed setting, you want to suppress the all-reduce communication for the first $K-1$ micro-batches and allow it only for the final micro-batch. This avoids $K-1$ unnecessary communication rounds that would each consume network bandwidth and introduce synchronization overhead.

To understand the cost savings, consider a training run with 8 GPUs and $K = 8$ accumulation steps. Without the no_sync() pattern, each backward pass triggers an all-reduce, for a total of 8 all-reduce operations per optimizer step. With no_sync(), only the final backward pass triggers the all-reduce, reducing communication to 1 operation per optimizer step. For models where communication is a significant fraction of step time (common for large models on fast GPUs connected by slower interconnects), this can yield substantial speedups.

PyTorch's DistributedDataParallel provides the no_sync() context manager for exactly this purpose. When a model is wrapped in no_sync(), gradients are accumulated locally on each GPU without triggering all-reduce. On the final micro-batch (outside no_sync()), the backward pass triggers the all-reduce as usual, averaging the accumulated gradients across all GPUs.