Gradient Accumulation › Gradient Accumulation for Memory

The memory savings from gradient accumulation are real and significant, but they apply only to a specific category of memory: activation memory. Understanding which types of memory are affected shows when accumulation helps.

GPU memory during training is occupied by several distinct components:

• Model parameters: The weights and biases themselves. For a model with $P$ parameters stored in float32, this costs $4P$ bytes. In mixed precision training, a float16 copy costs $2P$ bytes.
• Optimizer states: Adam maintains first and second moment estimates, costing roughly $8P$ bytes in float32 (two copies of the parameter count). This dominates for large models.
• Gradient buffers: One gradient tensor per parameter, typically matching the parameter precision. Costs $4P$ bytes in float32 or $2P$ in float16.
• Activation memory: Intermediate tensors computed during the forward pass and retained for the backward pass. Unlike the above, this scales with batch size. For a transformer, activation memory grows roughly as $O(L \cdot d \cdot B)$ where $L$ is the number of layers, $d$ is the model dimension, and $B$ is the batch size.

Gradient accumulation reduces only activation memory, since at any point only one micro-batch is being processed. Parameters, optimizer states, and gradient buffers remain constant regardless of $K$.

It is worth being specific about why activation memory scales with batch size. During the forward pass, the model computes intermediate representations at each layer: attention scores, attention outputs, feed-forward intermediate activations, and so on. These are not thrown away after the forward pass because the backward pass needs them to compute gradients via the chain rule. Each sample in the batch has its own set of these intermediate tensors. If you double the batch size, you double the number of samples in flight, and thus double the total activation storage. Gradient accumulation sidesteps this by processing one micro-batch at a time: the activations for micro-batch 1 are computed, used for the backward pass, and then freed before micro-batch 2 begins.

The gradient buffers, by contrast, are not freed between micro-batches. They must persist across the entire accumulation cycle so that gradients can be summed. The size of these buffers is proportional to the number of parameters, not the batch size. For a 125-million parameter model in float32, the gradient buffers occupy roughly 500 MB regardless of whether you accumulate 1 or 64 micro-batches.

The benefit of accumulation is largest when activation memory is the binding constraint. For very large models, the optimizer state alone can exceed available memory, and accumulation provides no relief. In that regime, techniques like mixed precision training (covered in the Training Infrastructure part), ZeRO optimizer sharding, and activation checkpointing are necessary.

For models where activation memory dominates, gradient accumulation is immediately effective. A model that fits only 2 samples per forward pass can use $K = 64$ to reach an effective batch size of 128, with the only cost being extra sequential forward-backward passes.

In practice, the combination of micro-batch size and accumulation steps is often the first knob to turn when a training run fails with an out-of-memory error. Before resorting to more complex techniques like activation checkpointing or model parallelism, halving the micro-batch size and doubling $K$ can free enough memory to continue training without any change to the effective batch size or optimizer behavior.