Gradient accumulation

Problem

A recipe says "batch size 256 sequences of 2,048 tokens" (524,288 tokens per optimizer step). The learning rate, warmup and weight decay in that recipe were tuned for gradients averaged over that many tokens. Your GPU holds activations for 4 sequences at a time. Shrinking the batch to 4 changes the noise in every gradient and, with it, the learning rate that works (see learning-rate schedules), so the recipe no longer applies.

Activation memory scales with the micro-batch. The gradient is just a sum over examples, and backpropagation in PyTorch adds into param.grad rather than overwriting it. So you can run the batch in pieces, let the gradients pile up, and step once. The trade is time for memory: the same arithmetic, done sequentially, with one step at the end. The questions are how to scale each piece so that the sum is the gradient you meant, and which other parts of the training loop quietly assume one forward pass per step.

Intuition

Think of the gradient as a weighted vote. The full-batch mean-loss gradient gives every token (or example) the same weight, 1/N where N is the total count. If you split the batch into micro-batches and average each one separately, every token in a micro-batch of 4 tokens gets weight 1/4 of that micro-batch's share, while every token in a micro-batch of 8 gets 1/8. Short micro-batches are over-represented. With fixed-size images that never happens: every micro-batch has the same count, and averaging the averages is the average. With language data, padding and ignore_index make the counts differ, and the error is silent.

The other trap is anything that runs per forward pass or per backward pass rather than per step. Batch normalization computes statistics per forward, so it sees micro-batches, not the batch. Clipping a micro-batch's gradient clips the wrong vector. DDP all-reduces after every backward, so with accumulation it would communicate k times per step unless you tell it not to.

Mechanics

The gradient is a sum

Let the batch have micro-batches B₁…B_k with n_j counted tokens each and N = Σn_j. The mean loss over the batch is

L = (1/N) Σ_j Σ_{t ∈ Bj} ℓ_t  ⇒  ∇L = Σ_j (1/N) ∇ Σ_{t ∈ Bj} ℓ_t

So each micro-batch should backpropagate its summed loss divided by the global N. The common recipe, loss = criterion(out, y) / k, backpropagates (1/k)(1/n_j)Σℓ_t. That equals the line above only when every n_j = N/k. CrossEntropyLoss with the default reduction='mean' divides by the number of non-ignored targets in its own call (ignore_index, reduction), which is exactly n_j. Use reduction='sum' and divide by N yourself.

N has to be known before the first backward. On one GPU, count the targets that are not ignore_index across all micro-batches up front (the data is already loaded, or the collate step can report it). Under DDP, ranks see different N; all-reduce the count first and scale by world size, because DDP averages gradients across ranks (the note on summed versus averaged losses in its docstring is the same issue across GPUs).

What the loop must do once per step

zero_grad (which sets .grad to None by default) runs once at the start of the step, not per micro-batch, or it erases what you accumulated. step, the learning-rate scheduler and any EMA of weights run once at the end. Clipping (clip_grad_norm_) runs once, after the last backward and before the step, on the accumulated gradient. With mixed precision, GradScaler scales each backward but unscales and checks for infinities once, before clipping. Step counters that drive warmup count optimizer steps, not micro-batches.

Normalization layers see micro-batches

_BatchNorm.forward normalizes with the current batch's mean and variance in training mode and updates the running estimates with momentum 0.1 on every forward call. Under accumulation both change: each micro-batch is normalized with its own, noisier statistics, and the running averages are updated k times per step, so the effective momentum per optimizer step is 1 − (1 − 0.1)^k. Neither is fixed by scaling the loss. The options are to accept it (micro-batches of 16+ are usually fine), to use SyncBatchNorm to pool statistics across GPUs (not across time), or to use a normalization that doesn't depend on the batch: LayerNorm and GroupNorm are identical under accumulation, which is one reason transformers are easy to accumulate.

DDP and no_sync

DistributedDataParallel hooks each parameter's gradient and, as buckets of 25 MiB fill during backward, launches an all-reduce that averages them across ranks. With accumulation that averaging would happen after every micro-batch, moving the full gradient over the network k times to produce the same sum. The no_sync context manager sets require_backward_grad_sync to False; the flag is read in _pre_forward, so the forward, not just the backward, must run inside the context. Gradients accumulate locally, and the last micro-batch runs outside it and syncs the sum once.

Worked example

A one-weight model ŷ = w·x + b at w = 0.5, b = 0, with per-token loss ½(ŷ − y)², stands in for a language model. The batch is 8 sequences of lengths 2, 6, 1, 3, 5, 2, 4, 1 (24 tokens), split into micro-batches of 2 sequences. Every number below comes from tested helpers in ml_models_docs.py.

Implementation

A DDP training step with token-count normalization, no_sync on all but the last micro-batch, and clipping, stepping and zeroing once:

import contextlib
import torch
import torch.distributed as dist
import torch.nn.functional as F

def train_step(model, optimizer, scheduler, micro_batches, max_norm=1.0):
    # micro_batches: list of (input_ids, labels), labels == -100 where padded
    local = sum((y != -100).sum() for _, y in micro_batches)
    total = local.clone()
    dist.all_reduce(total)                       # tokens across all ranks
    scale = dist.get_world_size() / total        # DDP averages over ranks

    optimizer.zero_grad(set_to_none=True)
    for i, (x, y) in enumerate(micro_batches):
        last = i == len(micro_batches) - 1
        ctx = contextlib.nullcontext() if last else model.no_sync()
        with ctx:                                # forward must be inside too
            logits = model(x)
            loss = F.cross_entropy(logits.flatten(0, 1), y.flatten(),
                                   ignore_index=-100, reduction="sum")
            (loss * scale).backward()
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
    optimizer.step()
    scheduler.step()                             # once per optimizer step
    return norm

With GradScaler, wrap each backward as scaler.scale(loss * scale).backward() and call scaler.unscale_(optimizer) once before clipping, then scaler.step and scaler.update. Under FSDP the same no_sync exists but keeps unsharded gradients on every rank, so it trades memory back for communication (named here; the distributed-training page will cover it).

Tradeoffs

• Time, not memory, is what you pay. Accumulating k micro-batches takes about as long as running the big batch would have, minus the GPU's loss of efficiency on smaller kernels (see GPU execution). If a larger micro-batch fits, prefer it.
• Activation checkpointing is the alternative. Recomputing activations in backward cuts memory per example instead of examples per pass; the two combine, and checkpointing is better when even one sequence is too large.
• Exactness needs the global count first. Token-count normalization requires knowing N (and, under DDP, an extra small all-reduce) before the first backward. Some codebases accept the per-micro-batch mean instead; with sequence packing to a fixed length the counts are nearly equal and the error is small.
• Batch statistics don't accumulate. BatchNorm and anything else computed per forward (in-batch negatives for contrastive losses, MoE load-balancing statistics) sees the micro-batch. For in-batch negatives that is a real loss of signal; gradient caching exists for that case.
• Dropout and data order are the same. Random masks per example and the sequence of examples don't change, so runs with and without accumulation differ only by floating-point summation order.
• Large effective batches still need retuning. Accumulation lets you reach a batch size; it doesn't make that batch size a good one. Past the critical batch size, more tokens per step buy little.

Connections

References

• Goyal et al. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. 2017.
• McCandlish, Kaplan, Amodei and the OpenAI Dota Team. An Empirical Model of Large-Batch Training. 2018.
• Li et al. PyTorch Distributed: Experiences on Accelerating Data Parallel Training. VLDB, 2020.
• Ioffe and Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML, 2015.
• Gao, Zhang, Han and Callan. Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup (gradient cache). RepL4NLP, 2021.
• PyTorch v2.15.0: DistributedDataParallel.no_sync, Optimizer.zero_grad and step, clip_grad_norm_, _BatchNorm.forward.