50. Gradient Accumulation for a Larger Effective Batch Size

Problem

Memory is often the limiting factor on how large a batch can fit on one device, but training dynamics can want a much larger effective batch than that. Gradient accumulation reconciles the two: run several small micro-batches forward/backward without stepping the optimizer, letting .backward() accumulate gradients into .grad across all of them, and only call optimizer.step() once at the end — simulating one large batch's gradient using several small batches' memory footprint.

Implement a training step that loops over a list of micro-batches, accumulates gradients correctly, and steps the optimizer exactly once.

Source: src/50_gradient_accumulation_loop.py

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

Examples:
>>> model = nn.Linear(2, 1)
>>> opt = torch.optim.SGD(model.parameters(), lr=0.1)
>>> loss_fn = nn.MSELoss()
>>> micro_batches = [(torch.randn(4, 2), torch.randn(4, 1)) for _ in range(2)]
>>> avg_loss = train_step_with_accumulation(model, micro_batches, loss_fn, opt, 2)
>>> isinstance(avg_loss, float)
True

Step-by-Step Approach

1. Call optimizer.zero_grad() exactly once, before the loop — not per micro-batch, or you'd erase the accumulation.
2. Loop over each (x, y) micro-batch, running the forward pass and computing loss_fn(model(x), y).
3. Divide that loss by accumulation_steps before calling .backward() — this is what makes the sum of per-micro-batch gradients equal the gradient of the mean loss over the full concatenated batch, for equal-size micro-batches.
4. Call .backward() on the scaled loss with no zero_grad() in between iterations, so gradients accumulate into each parameter's .grad.
5. Track the (already-divided) loss values and sum them for logging — this sum equals the average micro-batch loss.
6. After the loop, call optimizer.step() exactly once, using the fully accumulated gradient across all micro-batches.

The key insight is the math behind step 3: for K equal-size micro-batches with a mean-reduced loss, mean_loss_full_batch = (1/K) * Σ_k mean_loss_microbatch_k. Dividing each micro-batch's loss by K before backward reproduces the exact full-batch gradient, not an approximation of it — gradient accumulation with mean-reduced losses is mathematically exact, not a heuristic.

Reference solution

def train_step_with_accumulation(model, micro_batches, loss_fn, optimizer, accumulation_steps):
    optimizer.zero_grad()  # once per accumulation cycle, not per micro-batch
    total_loss = 0.0
    for x, y in micro_batches:
        pred = model(x)
        # dividing before backward makes the *sum* of these gradients equal
        # the gradient of the mean loss over the full concatenated batch
        loss = loss_fn(pred, y) / accumulation_steps
        loss.backward()  # accumulates into .grad; no zero_grad() in between
        total_loss += loss.item()
    optimizer.step()  # single step using the full accumulated gradient
    return total_loss  # already averaged since each term was pre-divided

Key Functions & Tricks

• optimizer.zero_grad() placement — called once outside the micro-batch loop, the single most common mistake to get backwards in this pattern.
• loss.backward() without an intervening zero_grad() — PyTorch's default behavior is to accumulate (add) into .grad, which is exactly what this pattern relies on.
• Pre-dividing the loss by accumulation_steps — the scaling correction that makes accumulated gradients mathematically equal a true full-batch mean-loss gradient.
• optimizer.step() called once, after the loop — the whole point is one parameter update per K micro-batches, not K updates.

How to Recognize This Pattern