40. Gradient Accumulation Across Microbatches

Problem

GPU memory caps the maximum batch size that fits in a single forward/backward pass, but the effective batch size you want for stable training is often larger, especially for long-sequence models where activations dominate memory. Gradient accumulation reconciles the two: split the target batch into smaller microbatches, run forward and backward on each one without stepping the optimizer, let .backward() accumulate gradients into .grad across calls, and only call optimizer.step() once after every microbatch is done.

Implement one full accumulation step. The trick is entirely in the loss scaling: get it wrong and the update silently behaves like a different learning rate.

Source: src/40_gradient_accumulation_loop.py

def train_step_with_accumulation(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    microbatches: list[tuple[torch.Tensor, torch.Tensor]],
    loss_fn,
) -> float:
    ...

Examples:
>>> model = torch.nn.Linear(2, 1)
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
>>> mbs = [(torch.randn(2, 2), torch.randn(2, 1)) for _ in range(4)]
>>> loss = train_step_with_accumulation(model, optimizer, mbs, torch.nn.functional.mse_loss)
>>> isinstance(loss, float)
True

Step-by-Step Approach

1. Call optimizer.zero_grad() exactly once, before the whole accumulation window — not per microbatch, or you'd throw away everything accumulated so far.
2. Loop over each microbatch and run a normal forward pass: preds = model(inputs).
3. Compute the microbatch's loss, then scale it down by 1 / num_microbatches before calling .backward(). Since .backward() adds into .grad rather than overwriting it, skipping this scale makes the accumulated gradient num_microbatches times too large.
4. Call .backward() on the scaled loss for each microbatch, without calling optimizer.step() in the loop.
5. Track total loss for logging by multiplying each scaled loss back up (loss.item() * n) so the reported number reflects the true average loss across the full effective batch, not the artificially shrunk per-microbatch value.
6. After the loop, call optimizer.step() exactly once — this applies a single update using the gradient accumulated across every microbatch, equivalent to having run the full batch at once.
7. Validate the equivalence directly: with equal-size microbatches and a mean-reduction loss, gradient accumulation must produce numerically the same parameter update as a single full-batch forward/backward/step.

The key insight is that .backward() is additive by default — that's precisely the mechanism gradient accumulation exploits, and precisely why the 1/n scale has to be applied before each microbatch's backward, not after the loop.

Reference solution

def train_step_with_accumulation(model, optimizer, microbatches, loss_fn):
    optimizer.zero_grad()   # clear once, before the whole accumulation window
    n = len(microbatches)
    total_loss = 0.0

    for inputs, targets in microbatches:
        preds = model(inputs)                  # forward on one microbatch only
        # Scale DOWN by 1/n before backward: .backward() *adds* into .grad,
        # so without this scale the accumulated gradient would be n times
        # the true full-batch-average gradient -- equivalent to silently
        # multiplying the learning rate by n.
        loss = loss_fn(preds, targets) / n
        loss.backward()                         # accumulates into .grad, no step yet
        total_loss += loss.item() * n           # unscale back for honest logging

    optimizer.step()                            # one update, using the fully accumulated gradient
    return total_loss / n

Key Functions & Tricks

• optimizer.zero_grad() — called once outside the microbatch loop, the most common placement bug in real accumulation code.
• loss / n before .backward() — the single line that makes accumulation mathematically equivalent to a full-batch update.
• .backward()'s additive accumulation into .grad — the underlying PyTorch behavior the whole technique relies on.
• loss.item() * n — unscaling for logging so displayed loss values don't silently shrink as microbatch count grows.
• copy.deepcopy(model) in the test harness — used to give the full-batch and accumulation paths identical starting weights for a fair comparison.

How to Recognize This Pattern