Gradient Accumulation

Gradient accumulation means running several forward and backward passes before a single optimizer step, so the gradients sum inside the parameter buffers and the update behaves as if you had used one large batch. The effective batch is the micro-batch size times the number of accumulation steps.

Say you want an effective batch of 512 because the loss curve is smoother and the optimizer step makes more sense at that scale. The accelerator on the desk holds 32 examples before it runs out of memory. Doubling the batch is not an option. Halving the model is not an option. The trick the field reached for in 2017 and kept using is to run 16 backward passes, let the gradients accumulate inside the parameter buffers, and only step the optimizer when the count reaches the target.

The risk is that the loss is no longer the same number it was at the bigger batch. The cross entropy of 16 mini-batches summed naively is 16 times the loss of one full batch. Without scaling, the gradient direction is correct but the magnitude is wrong, and the optimizer step is 16 times too big. The fix is one division, and it is easy to forget.

The contract has four parts:

• Divide each micro-batch's loss by accum_steps before backward(). PyTorch sums gradients into param.grad by default, and the division brings the running sum back to the right scale.
• Fire the optimizer step once per effective batch, after the last micro-batch's backward. Stepping mid-accumulation skews every parameter the rest of the run depends on.
• Advance the optimizer's state (momentum buffers, Adam moments) once per effective step, not once per micro-batch. Otherwise the exponential moving averages see the wrong frequency and burn through the schedule.
• On a single device this is bookkeeping. On a multi-rank cluster the same pattern wraps the non-final micro-batches in a no_sync context that skips the gradient all-reduce, so the last micro-batch reduces the full accumulated gradient in one pass instead of paying the network cost N times.

The equivalence proof in code

loss = criterion(model(x_full), y_full)
loss.backward()
opt.step()

is equivalent to

for x, y in chunks(x_full, y_full, n):
    scaled = criterion(model(x), y) / n
    scaled.backward()
opt.step()

up to floating point summation order. The accumulated gradient buffer at the end of the loop is the same tensor a single full-batch backward would produce. equivalence_check returns the max absolute gradient and parameter differences, and main() asserts both stay under 1e-4.

Where the cost goes

Each micro-batch costs one forward and one backward, so accumulation trades memory for time. The throughput curve in outputs/accum-curve.json shows what happens as the effective batch grows at fixed micro-batch:

There is no free lunch. Doubling accum_steps doubles the wall time per optimizer step. What changes is the variance of the gradient estimate: at the same wall budget you take fewer optimizer steps, but each one is averaged over more samples. The literature treats large batch and small batch as different optimization problems. The point here is mechanical, not statistical.

code/main.py is the runnable artifact. It does three things.

Step 1: equivalence check

equivalence_check() builds two copies of the same network with the same seed. One sees a 16-sample batch in one forward pass. The other sees four 4-sample chunks with the loss divided by four. The function compares the gradient buffers before the optimizer step and the parameters after. It returns max_grad_diff, max_param_diff and params_init_match, and main() asserts both diffs are under 1e-4.

Step 2: sync-on-last-step pattern

train_one_optimizer_step walks the micro-batches. Every micro-batch except the last enters no_sync_context(model). On a single process the context is a no-op; on DDP it is where the gradient all-reduce gets skipped. The bookkeeping is the same either way. A sync_counter increments on each synced backward (the last micro-batch, or every micro-batch if no_sync_until_last is False), so with the default it grows by one per effective step, not N, and accumulates across num_steps.

Step 3: the throughput curve

sweep_effective_batches runs the same model with a fixed micro-batch and a list of accumulation steps. For each setting it logs:

• samples_per_sec: total samples seen divided by wall time
• median_step_ms: 50th percentile per effective step
• sync_calls: collective points exercised
• avg_loss: average across the sweep's optimizer steps

The output lands in outputs/accum-curve.json and is reusable from a notebook.

Run it:

python3 code/main.py

The script prints the equivalence diff, then the sweep table, then the JSON path. Exit code zero.

In production training, gradient accumulation lives behind one knob. PyTorch's pattern is accumulation_steps = effective_batch // (micro_batch * world_size). That is the general formula for a multi-rank run; the demo here is a single process, so world_size is 1. The frameworks you are not allowed to use here wrap the same loop with the same steps: scale the loss, skip sync on non-final micros, accumulate, step once.

Three patterns show up in practice:

• Pick the micro-batch size to saturate device memory. Smaller wastes accelerator cycles, and larger crashes.
• Pick the effective batch from a learning rate schedule. Large effective batches need scaled learning rates and warmup, which is the linear scaling rule people have talked about since 2017.
• Treat the accumulation count as the bridge between the two. It is a knob you can tune at runtime without rewriting the data loader.

outputs/skill-gradient-accumulation.md captures the recipe so a peer can drop it into a new repo: scale loss by accum_steps, skip optimizer sync on non-final micros, step the optimizer once per effective batch, and log throughput against effective batch as JSON so the trade is visible.

1. Re-run the sweep with --num-steps 100 and plot samples per second against effective batch. Where does the curve flatten?
2. Add a wrong scaling variant (no division) and show the parameter diff at step 1 against the reference.
3. Swap SGD for AdamW and confirm the optimizer state advances once per effective step, not once per micro-batch.
4. Introduce a real DistributedDataParallel wrapper and route the no_sync_context to its method. Confirm the number of all-reduces per effective batch drops from N to 1, which is what sync_calls counts when no_sync_until_last is on.
5. Modify the equivalence check to compare two different micro splits (2 by 8 vs 4 by 4) and explain any tolerance you need to relax.

• PyTorch docs on DistributedDataParallel.no_sync for the production version of the sync-on-last-step trick.
• Goyal et al., 2017, on linear scaling for large batch training, the canonical reason to care about effective batch.
• PyTorch issue tracker on gradient accumulation interactions with mixed precision unscaling.
• Phase 19 lessons 42 to 45 cover the model, data loader, optimizer, and trainer scaffolding this lesson assumes.
• Phase 19 lesson 47 covers checkpoint and resume so a long accumulation run survives a wallclock cap.