Gradient accumulation in PyTorch: a tested training loop

Gradient accumulation combines gradients from several small batches before an optimizer updates a model. In PyTorch, you run a forward and backward pass for each microbatch, then call optimizer.step() once for the group. This lets you increase the number of examples contributing to an update while processing fewer examples at a time.

The details matter: the loss must use the right denominator, the final incomplete group needs an update, and gradients must stay intact between microbatches. This guide builds a small training loop that handles those cases and explains what changes with mixed precision and multiple GPUs.

How gradient accumulation works in PyTorch

PyTorch adds each backward pass's gradients to the parameters' existing .grad buffers. It does not automatically clear them. An accumulation loop uses that behavior deliberately: clear once at the start of an accumulation window, call backward() for each microbatch, then update the parameters once. The PyTorch guide to zeroing gradients explains the underlying behavior.

For equal-size microbatches, the nominal effective batch size is:

Microbatch size × accumulation steps × data-parallel workers.

For example, four examples per microbatch and three accumulation steps give 12 examples per update on one worker. With two data-parallel workers processing different examples at those same sizes, the nominal global batch is 24. These counts describe full windows; a shorter final window has fewer examples.

Keep the three counts separate in your run notes: microbatches processed, examples or valid tokens processed, and optimizer updates. Changing accumulation steps changes the number of updates per epoch unless another setting compensates.

Normalize the loss across the actual window

If every microbatch has the same number of equally weighted examples, dividing its mean loss by the number of microbatches in the window gives each example the same contribution to the update. However, dividing every loss by a fixed configured value fails when the last window is shorter.

Unequal microbatch sizes need another adjustment. A mean over four examples and a mean over one example must not receive equal weight. The example below sums each microbatch's losses and divides by the total number of examples in that window. It therefore handles both unequal microbatches and an incomplete final window.

This denominator is specific to ordinary classification with one class-index target per example, no class weights, and no ignored targets. For token training, count the valid prediction targets. With class weighting or ignored labels, follow your loss function's reduction rules; CrossEntropyLoss documents those distinctions. Do not substitute a sequence count for a token count without checking the objective.

A PyTorch gradient accumulation example

Save this as accumulation_demo.py in an environment with PyTorch installed. The standalone demo uses synthetic data and runs on CPU. Its data loader yields CPU tensors; the loop holds one window of input batches in host memory and transfers one microbatch at a time to the model's device.

from itertools import islice
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


def train_epoch(model, loader, optimizer, accumulation_steps=3):
    if accumulation_steps < 1:
        raise ValueError('accumulation_steps must be positive')
    model.train()
    device = next(model.parameters()).device
    criterion = nn.CrossEntropyLoss(reduction='sum')
    batches = iter(loader)
    update_sizes = []

    while True:
        window = list(islice(batches, accumulation_steps))
        if not window:
            break
        examples = sum(y.size(0) for _, y in window)
        optimizer.zero_grad(set_to_none=True)

        for x, y in window:
            x, y = x.to(device), y.to(device)
            loss = criterion(model(x), y) / examples
            loss.backward()

        optimizer.step()
        update_sizes.append(examples)
    return update_sizes


if __name__ == '__main__':
    torch.manual_seed(7)
    x = torch.randn(19, 5)
    y = torch.randint(0, 3, (19,))
    loader = DataLoader(TensorDataset(x, y), batch_size=4,
                        shuffle=False, drop_last=False)
    model = nn.Linear(5, 3)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    print(train_epoch(model, loader, optimizer))

Run it with:

python accumulation_demo.py

The output lists how many examples contributed to each optimizer update:

[12, 7]

The first update uses three microbatches of four examples. The second uses two microbatches, containing four and three examples. Both losses use their window's actual example count, so the last seven examples are included with the intended mean-loss weighting.

This grouping approach trades some host memory for an explicit denominator. Keep loader outputs on CPU and account for the size of the buffered window. For streaming or unusually large inputs, design a bounded grouping strategy that preserves the same normalization. Preloading the whole window onto the GPU would undermine the intended memory benefit.

What the example tests establish

We ran the demo with PyTorch 2.15.0+cpu. Ten additional comparisons tested SGD with momentum and AdamW across full windows, incomplete windows, unequal microbatches, a dataset smaller than one microbatch, and one accumulation step. In a deterministic network without batch-dependent or stochastic layers, accumulated gradients and updated parameters matched concatenated-batch calculations at every update within the specified numerical tolerance.

Those comparisons used CPU float64 tensors. They verify the loop's arithmetic and update boundaries; they do not measure GPU memory savings, throughput, mixed-precision behavior, or distributed execution.

Preserve the optimizer update boundary

Calling zero_grad() inside the microbatch loop would erase earlier contributions. Calling step() there would change the parameters before the window is complete. Both changes defeat this example's intended comparison with one larger-batch update.

Likewise, do not add microbatch loss tensors together and postpone all backward passes if the aim is to release each microbatch's saved activations promptly. The example calls backward separately. Ordinary accumulation across fresh forward passes does not require retain_graph=True; see the backward API's graph-lifetime rules.

Define a learning-rate scheduler's unit before attaching it to the loop. If your schedule is defined per optimizer update, advance it at that boundary rather than after every microbatch. If it is defined per epoch or validation event, preserve that meaning. Record the resulting update count so a changed accumulation setting does not silently alter warmup or decay.

Mixed precision and gradient clipping

For FP16 training with a gradient scaler, keep gradients scaled and the scale unchanged throughout the window. After the last backward pass, unscale once if clipping is needed, clip the accumulated gradients, then call the scaler's step and update methods. Nonfinite gradients can cause the optimizer step to be skipped. Follow the official AMP accumulation example when adapting this loop.

For device setup and precision choices, use our PyTorch on A100 guide. The CPU example here does not validate a CUDA or AMP configuration.

Multiple GPUs and DDP synchronization

With PyTorch DistributedDataParallel, no_sync() can defer gradient synchronization for the earlier microbatches. Include both forward and backward inside that context, then run the final pair outside it. Synchronize before the optimizer update. The DDP documentation specifies this behavior.

The simple global-batch formula assumes equal per-worker contributions. Unequal sample or valid-token counts need normalization across workers, accounting for DDP's gradient averaging. Coordinate window boundaries across ranks. The single-process example is not a drop-in distributed training loop.

When accumulation differs from a larger physical batch

Our equivalence test deliberately excludes layers that couple examples within a batch. Batch normalization computes statistics from the current minibatch, so splitting that batch can change both outputs and running statistics. Dropout and other random operations can also produce a different numerical trajectory. A matching effective batch size does not guarantee matching training results.

Compare runs with controlled data order, initial weights, precision, optimizer settings, and update counts. PyTorch's reproducibility guidance also cautions that identical seeds do not guarantee identical results across platforms or releases. Use tolerances and validation metrics appropriate to the workload.

Accumulation reduces the activation burden of processing a large physical batch, but it does not remove the model's parameters, gradients, or optimizer state. It also adds repeated forward and backward work per update. Measure memory and examples or valid tokens per second before choosing a setting. More accumulation steps are not a promise of faster training.

Checks before a longer training run

• Confirm every intended example reaches an optimizer update, including the final partial window.
• Check loss weighting against a small concatenated-batch reference.
• Log optimizer updates separately from microbatches and data volume.
• Clip and schedule at the boundaries your training design requires.
• Compare validation quality alongside memory use and throughput.
• Save checkpoints at completed update boundaries unless your implementation also saves unfinished gradient state.

If you use a trainer that already implements accumulation, verify its settings before adding a second accumulation loop. Our LoRA fine-tuning guide shows where that setting appears in a larger workflow. Start with the smallest controlled run that proves the update behavior you need.