Fill the conditions that preserve accumulation to large batch equivalence
Accumulation equals one large batch only if token losses are divided by the whole step's token count, the optimizer steps once per full window, and loss scaling and precision stay fixed across micro-batches.
A head teacher wants the average test score of every student in four classrooms. One quick way is to average the four class averages. That works only if each room has the same number of students. If one room has 5 students and another has 50, the small room counts as much as the big one, and the result is wrong. The right way is to add up every score and divide once by the total number of students. The teacher must also wait until all four rooms report, and every room must mark on the same scale. Gradient accumulation is the same: add every token's loss, divide by all the tokens, update once, and keep the scale fixed.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Gradient accumulation exists because of memory. A GPU may hold 4 sequences at once when the training recipe calls for 16 per GPU. So the training loop runs 4 micro-batches, adds their gradients, and takes one optimizer step. The goal is a step that is the same as one step on all 16 sequences.
It can be exactly the same, apart from floating-point rounding. Three conditions decide it: the loss is normalized the same way over the whole step, the optimizer steps after the same amount of data every time, and loss scaling and precision are handled the same way across the micro-batches.
Why gradients can be added
The loss for one step of next-token training is the mean cross-entropy over all loss tokens in the batch. A loss token is a position whose label counts, so padding and masked prompt tokens are left out. If the batch has N loss tokens, the gradient is:
Differentiation is linear, so this sum can be split into any groups of tokens. Split the batch into k micro-batches, compute the sum of token gradients in each, and add the k results. Divide by N once and you have g exactly. The order of the additions changes the floating-point rounding, but nothing else.
That argument needs a loss that is a sum over tokens, with no part that depends on the batch as a whole. Next-token cross-entropy meets that need. LayerNorm and RMSNorm work per token, so they do not care how the batch is split. Three common pieces do not meet it:
- BatchNorm computes statistics over the batch. LLMs rarely use it, but many vision models do.
- A contrastive loss with in-batch negatives compares each example with the others in its batch, so a micro-batch sees fewer negatives.
- An MoE load-balancing loss uses the share of tokens sent to each expert over the batch. Computed per micro-batch, it enforces balance on a smaller sample than intended.
For these, accumulation gives a different objective, not only a slower one.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
from contextlib import nullcontext
from itertools import batched # Python 3.12+
import torch, torch.nn.functional as F, torch.distributed as dist
accum_steps = 4
world = dist.get_world_size() if dist.is_initialized() else 1
optimizer.zero_grad(set_to_none=True)
for window in batched(loader, accum_steps): # one optimizer step
# Loss tokens in the whole step, summed over all ranks.
n_tok = sum((b["labels"] != -100).sum() for b in window)
if world > 1:
dist.all_reduce(n_tok)
for i, b in enumerate(window):
last = i == len(window) - 1
with nullcontext() if last or world == 1 else model.no_sync():
with torch.autocast("cuda", dtype=torch.bfloat16):
logits = model(b["input_ids"]).logits
loss_sum = F.cross_entropy(
logits.flatten(0, 1).float(), b["labels"].flatten(),
ignore_index=-100, reduction="sum") # labels pre-shifted
# DDP averages over ranks, so scale by world to get sum / n_tok.
(loss_sum * world / n_tok).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.step()
optimizer.zero_grad(set_to_none=True)
Real products, models, and research that use this idea.
- Hugging Face Transformers fixed its Trainer in October 2024. The fix divides the accumulated loss by the total non-padding tokens of the whole accumulation step, passed in as num_items_in_batch, instead of averaging per-batch means.
- With PyTorch AMP, an accumulation loop keeps one loss scale for every micro-batch of an effective batch. It calls unscale_ only after all gradients are accumulated, and calls scaler.update() only after a full step.
What an interviewer would ask next. Try answering before peeking at the approach.
QHow can you test that an accumulation setup matches the large batch before a long run?
From one checkpoint, compute one step with k = 1 and one with k = 4 at a quarter of the micro-batch, on the same data, and compare the gradients.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Dividing each micro-batch's mean loss by the accumulation count. That is exact only when every micro-batch has the same number of loss tokens; with padding or packing, short micro-batches get too much weight.
60 second bullets to scan on the way to the call.
Why accumulated gradients can equal the large-batch gradient
Mean of micro-batch means versus one token-weighted mean
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.