Gradient Accumulation

1. What Happens in a Training Update

The model's parameters, often called weights, are the numbers training adjusts. For a language model, a training example is usually a sequence of tokens, the small pieces into which text is split. During the forward pass, the model uses its current weights to predict the next token at each eligible position.

A loss measures how well those predictions match the targets. Giving the correct next token a low probability produces a larger loss. The backward pass, or backpropagation, then computes the loss's gradient with respect to each trainable parameter. A gradient tells us how a small change in that parameter would affect the loss near its current value.

An optimizer, such as SGD or AdamW, uses the gradients to update the weights. The learning rate controls the scale of the update. Computing a gradient and applying an update are separate operations: a backward pass can finish while the weights remain exactly where they were.

This separation is visible in PyTorch. Calling loss.backward() adds gradients to the parameters' .grad fields, while optimizer.step() applies an update. optimizer.zero_grad() clears the accumulated gradients so the next update can start fresh. PyTorch's gradient-clearing guide explains this default accumulation behavior.

2. Why Larger Batches Need More Memory

A batch groups training examples together. Instead of basing an update on one sequence, we might use eight, thirty-two, or more. Their losses contribute to a combined objective, usually an average over the prediction targets included in the batch.

The weights are shared across the examples, but the intermediate results of the forward pass depend on each example's text. These intermediate tensors are called activations. Many must remain available until the backward pass uses them. Processing more sequences together generally increases this activation memory, especially when the sequences are long.

Suppose eight sequences should contribute to one update, but only two fit comfortably in one forward-and-backward pass. Updating the weights after every pair would fit in memory, but would produce four small-batch updates. Gradient accumulation instead keeps the intended single update and changes how we schedule the calculations leading up to it.

3. Splitting a Batch into Smaller Passes

Each smaller group is a microbatch. A set of microbatches whose gradients contribute to one optimizer update is an accumulation window. In our example, a microbatch contains two sequences and the window contains four microbatches, giving an effective batch size of eight sequences.

Start with an empty gradient buffer. Run the first microbatch forward and backward, keeping its gradients. Process the next microbatch using the same weights and add its gradients to the buffer. After all four have contributed, apply one optimizer update and clear the buffer for the next window.

For this first example, all four microbatches contain the same number of prediction targets. We divide each microbatch's mean loss by four before its backward pass, so the accumulated gradient represents the average over the whole window. The visual follows one weight to make the arithmetic visible.

Once a microbatch's backward pass finishes, its saved activations can be released. Its contribution remains in the gradient buffer, which is added to rather than stored as a separate copy for every microbatch. Ordinary accumulation does not require retain_graph=True. Keeping all the forward-pass graphs and calling backward only at the end would retain the activation memory we wanted to avoid. The backward documentation describes when the graph is freed.

4. Why the Gradients Need Averaging

The gradient of a sum of losses is the sum of their gradients. An average works the same way, with a division by the number of equally weighted terms. This lets us compute the pieces separately and still obtain the gradient of the combined loss, provided every piece uses the same model weights and the same weighting as the full batch.

The visual's four mean gradients are 2, −1, 3, and 4. Their average is 2. Simply adding those means would give 8, a gradient four times as large for this objective. Dividing each contribution by four gives 0.50, −0.25, 0.75, and 1.00, which accumulate to the intended value of 2.

On one device, equal-size microbatches give an effective batch size of microbatch size times accumulation steps. Increasing the number of accumulation steps increases the amount of data per update, not the number of examples processed simultaneously. This formula counts sequences; the number of valid token targets also depends on their lengths and label masks.

5. For Language Models, Count the Right Tokens

Two text sequences do not necessarily supply the same number of training targets. One may be shorter, and padding positions should not contribute to its loss. Instruction tuning can also exclude parts of the prompt so that only selected targets are supervised. Here, a valid target means a token prediction included in the loss.

For the usual token-mean language-model objective, each valid target should have equal weight across the entire accumulation window. Add their individual losses and divide by the total number of valid targets. Averaging microbatch means equally is correct only when those microbatches contain equal numbers of contributing targets. Hugging Face's gradient accumulation explanation discusses why this distinction matters.

In the unequal case, the first microbatch has a loss sum of 2 and the second has a loss sum of 18. The combined mean is 20 divided by 8, or 2.50. Backpropagating the first sum divided by 8 and then the second sum divided by 8 produces the correctly weighted gradient. There is no additional division by the number of microbatches.

In next-token training, the prediction at one position is matched with the label at the following position. Count valid targets after this shift, excluding ignored labels. With the conventional unshifted label tensor, that means counting over labels[:, 1:]. The current Transformers guide makes this denominator explicit. A deliberately sequence-weighted objective would require different weighting; the examples here use a token mean.

6. A PyTorch Loop That Preserves the Token Mean

The example below groups incoming microbatches into windows. Before processing a window, it counts the valid targets across that window on the CPU. It then moves and processes one microbatch at a time, dividing every summed loss by that same target count. The optimizer runs only after the window is complete.

This is a single-device, full-precision example for a causal language model that returns .logits. The model must already be on device. The loader supplies CPU tensors named input_ids, attention_mask, and labels; labels are unshifted, with padding and other excluded targets set to -100. The example optimizes only the next-token cross-entropy loss.

from itertools import islice

import torch
import torch.nn.functional as F


def train_epoch(model, loader, optimizer, device, accumulation_steps=4):
    if accumulation_steps < 1:
        raise ValueError("accumulation_steps must be positive")

    model.train()
    optimizer.zero_grad(set_to_none=True)
    batches = iter(loader)
    updates = 0

    while True:
        # Keep this window of input batches on the CPU.
        window = list(islice(batches, accumulation_steps))
        if not window:
            break

        counts = [
            int(batch["labels"][:, 1:].ne(-100).sum())
            for batch in window
        ]
        target_count = sum(counts)
        if target_count == 0:
            continue

        for batch, count in zip(window, counts):
            if count == 0:
                continue

            inputs = {
                name: tensor.to(device)
                for name, tensor in batch.items()
                if name != "labels"
            }
            labels = batch["labels"].to(device)
            logits = model(**inputs, use_cache=False).logits
            loss_sum = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
                reduction="sum",
            )
            (loss_sum / target_count).backward()
            del inputs, labels, logits, loss_sum

        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        updates += 1

    return updates

The last window may contain fewer than four microbatches. It still receives an update, using its actual target count rather than the count of a hypothetical full window. A microbatch with no valid targets contributes nothing, and an entirely empty-target window is skipped. Only input tensors are held for the window; activation graphs are handled one microbatch at a time.

The code computes the loss explicitly, so it does not also use a model-provided mean loss. Training frameworks may manage accumulation and normalization for you. For example, Accelerate provides an accumulation context; adding manual division on top of a framework's existing normalization can scale the gradients twice.

7. What Belongs at the Update Boundary

Gradient clipping limits the size of an unusually large gradient. If the goal is one effective-batch update, clip the combined, correctly normalized gradient after all microbatches have contributed. Clipping each microbatch separately can change the final direction because clipping is not a linear operation.

AdamW's optimizer state and weight decay should advance with the optimizer update, rather than with each backward pass. A learning-rate schedule defined in optimizer steps should likewise advance after an actual update. Four microbatches in one window count as one update for that schedule.

8. Accumulation Across Multiple GPUs

In data parallelism, several workers train copies of the same model on different examples and combine their gradients. With equal-size microbatches and matching accumulation windows, the global effective batch counts the sequences from every data-parallel worker.

Count data-parallel workers here, not every GPU in a model-parallel setup. GPUs that jointly process the same examples do not multiply the effective batch in this way. For variable-length text, normalization must account for valid targets across both the accumulation window and the data-parallel workers.

PyTorch DistributedDataParallel normally synchronizes gradients during backward. Its no_sync() context can defer that communication on the earlier microbatches, with synchronization on the final backward pass. Both the forward and backward calls belong inside the context for a microbatch that should not synchronize.

9. What Changes in the Training Run

With fixed weights, correct normalization, and a loss that separates across examples or target tokens, accumulation gives the same mathematical gradient as processing the effective batch together. It does not promise bit-for-bit identical results: floating-point addition order and random operations such as dropout can differ. PyTorch's numerical-accuracy notes explain why changing how calculations are grouped can alter rounding.

Layers that couple examples introduce another distinction. For example, BatchNorm computes training statistics from the batch it sees, so several small passes do not reproduce the statistics of one large pass. Standard decoder-only Transformers commonly use normalization within each token's representation, but the equivalence still depends on the actual model and loss.

Holding the effective batch fixed while splitting it into smaller microbatches mainly changes memory use and execution. Holding the microbatch fixed while increasing accumulation steps changes the training schedule: more examples contribute to each update and there are fewer updates for the same amount of data. Learning-rate and warmup settings should be considered in that context, rather than automatically multiplied by the accumulation factor.

The memory reduction mainly comes from processing fewer examples' activations at once. Model weights, optimizer state, and persistent gradient buffers still need space, and a single long sequence must still fit. LoRA, activation checkpointing, and sharding address other parts of the memory budget and can be combined with accumulation.

Accumulation also runs several smaller passes sequentially, so it is not a general throughput improvement over a larger physical batch that already fits. Its role is to separate the amount of data used for one update from the amount processed at one time: finish each microbatch's backward pass, keep its properly weighted gradient contribution, and update the model once the intended batch is complete.