Gradient Accumulation › Combining Accumulation with Other Memory Techniques

Gradient accumulation is rarely used in isolation. It combines naturally with other memory-saving techniques, and understanding how they interact helps you design an efficient training configuration.

Activation checkpointing (also called gradient checkpointing) reduces activation memory by recomputing some activations during the backward pass instead of storing them from the forward pass. Where gradient accumulation reduces activation memory proportional to micro-batch size, activation checkpointing reduces it proportional to the number of checkpointed layers, at the cost of one additional forward pass per backward pass.

The two techniques are orthogonal: you can use both simultaneously. With activation checkpointing, even a micro-batch of size 8 can have dramatically reduced activation memory, which means you can potentially increase $B_{\text{micro}}$ before resorting to accumulation. The tradeoff is that checkpointing increases computation time (roughly by 33%, since you recompute one forward pass), while accumulation does not increase the total computation per token but does add kernel launch overhead.

For large models where even single-sample activation memory is significant, activation checkpointing allows you to use meaningful micro-batch sizes that maintain good GPU utilization. Combining both techniques (checkpointing plus accumulation) gives you control over both computation overhead and communication overhead.

The ZeRO (Zero Redundancy Optimizer) family of techniques shards optimizer states, gradients, and parameters across GPUs to reduce per-GPU memory. ZeRO-1 shards optimizer states, ZeRO-2 additionally shards gradients, and ZeRO-3 shards all three including parameters.

When using ZeRO-2 or ZeRO-3, gradient accumulation still works but interacts with the sharding pattern. With ZeRO-2, gradients are gathered and reduced across GPUs at each backward pass. The no_sync() pattern for suppressing all-reduce during accumulation needs to be coordinated with ZeRO's gradient reduction hooks. Libraries like DeepSpeed handle this automatically through their gradient accumulation integration.

The combination of ZeRO-3 and gradient accumulation is particularly powerful for very large models: ZeRO-3 eliminates the need to hold full parameter, gradient, and optimizer state copies on each GPU, while gradient accumulation controls the activation memory and allows matching the effective batch size to the training recipe.

Mixed precision training, as discussed in the context of the mixed precision section above, reduces parameter and activation memory by using float16 or bfloat16 for most computations. The main-copy parameters are maintained in float32 for numerical stability.

Mixed precision roughly halves activation memory (since float16 activations are half the size of float32), which means you can often double the micro-batch size. If doubling the micro-batch is enough to reach the target effective batch size without accumulation, you can potentially eliminate accumulation entirely, trading the overhead of $K$ kernel launches for the overhead of the mixed precision scaler.

When both techniques are used together, the combined effect is multiplicative: mixed precision halves activation memory, and gradient accumulation by a factor of $K$ further reduces the effective activation footprint to $1/(2K)$ of the full-batch float32 baseline.