Gradient Accumulation for Larger Batches

The Standard Training Loop and Its Memory Bottleneck

Let's first revisit the standard training loop for a single batch, which we've been implicitly assuming until now. The process is a tightly coupled sequence:

1. Forward Pass: The model processes the batch of data to generate predictions.
2. Compute Loss: The predictions are compared to the true labels to calculate the loss.
3. Backward Pass: Gradients of the loss with respect to all model parameters are calculated via backpropagation (loss.backward()).
4. Update Parameters: The optimizer uses the calculated gradients to update the model's parameters (optimizer.step()).
5. Reset Gradients: The gradients are cleared to prepare for the next batch (optimizer.zero_grad()).

This entire cycle repeats for every batch in the dataset. The critical point is that the memory required for the intermediate activations in step 1, which are needed for backpropagation in step 3, scales directly with the batch size. This is often the primary memory bottleneck.

The flowchart on the left in the following image illustrates this standard, sequential process.

The Core Principle: Decoupling Updates from Gradients

Gradient accumulation's central innovation is to break the rigid link between the backward pass (step 3) and the parameter update (step 4). It exploits a key feature of modern deep learning frameworks like PyTorch: the calculation, application, and clearing of gradients are distinct, controllable operations.

When loss.backward() is called, the framework computes gradients and adds them to a .grad attribute on each model parameter. If you call loss.backward() again before clearing the gradients, the new gradients are simply summed with the existing ones. The optimizer only uses these accumulated gradients when optimizer.step() is finally called, and they are only reset to zero upon an explicit optimizer.zero_grad() call.

Gradient accumulation leverages this behavior to simulate a large batch. Instead of one large batch, we process several smaller micro-batches that fit into memory. We let the gradients from these micro-batches accumulate and only perform the optimizer step and zeroing after a specified number of micro-batches have been processed.

To understand this process, let's turn to a detailed explanation.

The flowchart on the right in the image we saw earlier perfectly visualizes this modified loop. We iterate through several mini-batches, accumulating gradients after each backward pass. Only when the "Accumulation Step" is reached do we update the parameters and reset the gradients, completing one step for the full effective batch.

The Mathematical Foundation and Its Practical Limits

In an ideal scenario, this technique is mathematically equivalent to training on a single large batch. Let $B$ be the large effective batch, which is split into $N$ micro-batches $b_1,b_2,\dots ,b_N$. The standard update rule using the large batch is:

where $w$ are the model weights, $\eta$ is the learning rate, and ${\nabla }_w\mathcal{L}(B)$ is the gradient of the loss for the entire batch $B$.

With gradient accumulation, we perform the update after accumulating gradients from all micro-batches:

If the total loss is simply the sum of the micro-batch losses (i.e., $\mathcal{L}(B)={\sum }_{i=1}^N\mathcal{L}(b_i)$), then by the linearity of the differentiation operator, these two updates are identical.

However, this theoretical equivalence is fragile. There are two critical situations where it breaks down, which you must be aware of to correctly implement and debug large-scale training.

1. The Batch Normalization Conflict

The most well-known issue is the incompatibility with standard Batch Normalization (BN) layers. BN normalizes a layer's activations using the mean and standard deviation computed across the current batch. When using gradient accumulation, BN computes these statistics over the small micro-batch, not the large effective batch. This creates a severe mismatch: the forward pass statistics are noisy and high-variance, while the backward pass gradient update corresponds to the larger, more stable effective batch. This desynchronization can destabilize training entirely.

This is a primary reason why Transformers do not use Batch Normalization. Instead, they use Layer Normalization, which computes normalization statistics across the feature dimension for each sample independently. Since Layer Normalization is batch-size-agnostic, it is perfectly compatible with gradient accumulation.

2. The Loss Normalization Bug in LLMs

A more subtle but equally critical issue arises in how loss is typically calculated for language models. The cross-entropy loss is often averaged over the number of non-padded tokens in a batch to ensure that sequences of different lengths contribute fairly.

When using gradient accumulation, a naive implementation might average the loss for each micro-batch independently and then simply sum the gradients. This is incorrect. If each micro-batch has a different number of tokens (due to variable sequence lengths and padding), this process does not produce the same gradient as if you had computed the loss over the total number of tokens in the entire effective batch at once. This bug was discovered to be present in many popular libraries.

The following article explains this issue, which is highly relevant for anyone wanting to train LLMs from scratch.

The correct approach is to sum the losses from each micro-batch without averaging them, and then normalize this total loss by the total number of tokens across all accumulated micro-batches before performing the backward pass. Modern libraries like Hugging Face's transformers now implement this corrected logic.

Strategic Advantages and The Final Verdict

Despite these nuances, when implemented correctly, gradient accumulation is a powerful tool with clear advantages.

The core takeaway is that gradient accumulation is not a free lunch. It is a direct trade-off: you save memory at the cost of increased training time because the multiple forward/backward passes on micro-batches are processed sequentially, incurring more overhead than a single, highly parallelized pass on a large batch. Therefore, the rule is simple: if your desired batch size fits comfortably in memory, do not use gradient accumulation. It is a tool of necessity, not a general performance optimizer.

Conclusion

In this lesson, we have demystified gradient accumulation, a critical technique for training large-scale models. By understanding its mechanism, we can now see how to train with the large, stabilizing batch sizes that modern optimizers and learning rate schedules demand, even when faced with hardware memory constraints.

Here are the key takeaways:

• Core Mechanism: Gradient accumulation simulates a large effective batch by running forward and backward passes on several smaller micro-batches and summing their gradients before performing a single optimizer step.
• The Trade-off: It trades GPU memory (a spatial cost) for longer training time (a temporal cost).
• Mathematical Equivalence: While theoretically sound, the equivalence to true large-batch training can be broken by practical details, most notably the use of Batch Normalization and incorrect loss scaling with variable-length sequences.
• Primary Use Case: Its purpose is singular: to enable training with batch sizes that would otherwise cause out-of-memory errors. It should not be used if memory is not a constraint.

We've now added another essential tool to our pre-training toolkit. Gradient accumulation tackles the memory cost of large batches. In our next lesson, we will explore mixed-precision training, another synergistic technique that reduces the memory footprint of the model's weights, activations, and gradients, often leading to significant speedups in the process.