Gradient Accumulation

Large batches are wanted for stable gradients and good hardware utilisation. They are limited by activation memory, which scales linearly with batch size.

The gradient of a sum is the sum of gradients, so the batch can be split.

Run each micro-batch forward and backward, let the gradients accumulate in .grad, and step after the last one. Activations for one micro-batch at a time, optimiser behaviour of the full batch.

The normalisation bug

Eq. 1 divides by $|\mathcal{B}|$ — the total number of examples, or for a language model the total number of tokens. The natural implementation divides each micro-batch’s loss by the number of accumulation steps instead.

Those agree only when every micro-batch has the same token count.

The second is a mean of means. A micro-batch of 100 tokens and one of 1000 contribute equally, so short sequences are weighted ten times too heavily.

This was found in late 2024 to affect several widely used fine-tuning libraries. The symptom was that runs with accumulation reached measurably worse loss than the equivalent large batch — attributed for some time to accumulation being “approximate”, which it is not. It is exact, and the implementations were wrong.

import torch
from torch import nn


def train_step(model: nn.Module, opt, batches: list, clip: float = 1.0) -> float:
    """One optimiser step over several micro-batches."""
    opt.zero_grad(set_to_none=True)

    # Count the tokens first: the denominator is the total, not the batch count.
    total_tokens = sum(int(b["labels"].ne(-100).sum()) for b in batches)
    running = 0.0

    for i, b in enumerate(batches):
        # Sum reduction, scaled once by the global denominator.
        out = model(**b, reduction="sum")
        loss = out.loss / total_tokens

        # Skip the all-reduce on every micro-batch but the last.
        sync = i == len(batches) - 1
        with model.no_sync() if not sync and hasattr(model, "no_sync") else nullcontext():
            loss.backward()

        running += loss.item()

    # Clip after accumulation, on the complete gradient.
    torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
    opt.step()
    return running

Two details beyond the denominator.

no_sync matters under data parallelism. Without it, DDP all-reduces the gradient after every micro-batch — sixteen collectives per step where one is needed, which on a slow interconnect is most of the step time.

And clipping belongs after the loop. Clipping each micro-batch bounds a partial gradient and changes the direction of the accumulated one, which is the property clipping exists to preserve.

Batch normalisation, if any survives in the model, computes statistics per forward pass — so a micro-batch of 4 gives statistics of 4, not of 64. Accumulation is exact for everything else and not for that, which is one more entry on BatchNorm’s list of incompatibilities.

The hardware limit also remains. Micro-batches below about 8 leave the GPU underutilised, so accumulation trades memory for wall-clock beyond that point rather than getting it free.