Train with a batch that does not fit, by not holding it all at once. The gradient is a sum, and sums can be built up.
Overview
Why it works
The loss over a batch is the mean of the per-example losses, and differentiation is linear. So the grad of the batch is the mean of the per-example grads.
That means the sum can be built in pieces. Run a forward and backward pass on 16 examples, keep the grads, run another 16, add to them, and after 16 rounds you've just the grad of a 256-example batch. Step the tuner once.
The resulting update is the same. Not similar, the same up to floating-point sorting — to what a single 256-example batch would have made.
Gradient Accumulation
This module needs JavaScript: the
numbers are computed in the page rather than recorded.
Worth knowing
The gradient of a mean loss over a batch is the mean of the per-example gradients. So it can be accumulated in pieces.
Run several small forward and backward passes, sum the gradients, and step once. The update is identical to the large batch.
Only activation memory scales with the micro-batch. Weights, gradients and optimiser state do not.
It costs wall-clock time, not correctness — except for batch normalisation, which sees only the micro-batch.
Gradient Accumulation
Getting the effect of a large batch out of a device that can't hold one.
What actually consumes the memory
The bars separate the two cases at the memory budget you choose.
Training memory is roughly:
Fixed costs. Weights, their grads, and the tuner state. Adam keeps two extra tensors per param. So This is around four times the model size, and it doesn't depend on the batch at all.
Activations. each mid value from the forward pass has to be kept for the backward pass. This scales linearly with the batch size, and for a deep model it dominates.
Buildup attacks only the second. This Is why it works. Drag the micro-batch down and watch the second bar fall while the fixed part stays put.
The dashed line is the budget. When the full batch crosses it and the micro-batch doesn't, that gap is the entire reason the technique exists.
What it costs
Time. Sixteen sequential passes take about if sixteen passes. Buildup doesn't make any of it faster. It makes a thing possible.
There's a small speed loss too, since very small micro-batches use the accelerator less well. A micro-batch of 1 is often much less than a sixteenth as fast as a micro-batch of 16.
The one thing that is not identical
Batch norming. Its statistics are found over what is in the forward pass, which is the micro-batch, not the built up total.
So a model with batch norm trained with a micro-batch of 4 behaves like one trained with a batch of 4, but many buildup steps follow. The norming is noisier, and at very small micro-batches it degrades badly.
This is a real reason to prefer layer norming or group norming, both of which norm per example and are indifferent to the batch. It isn't a coincidence that transformers use layer norm and are the architectures trained with the most aggressive buildup.
Getting it right
Two details cause most of the bugs.
Divide the loss by the number of buildup steps before the backward pass, or the built up grad is a sum where the tuner expects a mean, and your effective learning rate is times'd by the step count.
Zero the grads at the right time. After the tuner step, not after each backward pass. Zeroing each pass defeats the whole thing, and it fails silently: training still runs, just with the small batch you were trying to avoid.
Alongside
Activation checkpointing trades find for memory in the other direction: discard activations during the forward pass and recompute them in the backward pass, for roughly 30% more time and a large memory saving. The two combine.
Mixed precision halves activation memory outright, and mixes with both.
Model parallelism and ZeRO attack the fixed costs, which buildup can't touch. When the weights alone don't fit, buildup isn't the answer.
Where the memory actually goes
Buildup works because it attacks the one cost that scales with the batch. Training memory splits in two: a fixed part. Weights, their grads, and the tuner state, which for Adam is about four times the model size — and activations, each mid value kept for the backward pass, which grows linearly with the batch. Take a 1 GB model on an 8 GB budget:
batch
fixed (Adam ~4×)
activations
total
256 (wanted)
4.0 GB
~5.1 GB
9.1 GB — over budget
16 (micro-batch)
4.0 GB
~0.3 GB
4.3 GB — fits
Run the micro-batch of 16 sixteen times, summing the grads, and step once: the update is *the same* to the 256-batch. The batch grad is the mean of the per-example grads and a mean can be built in pieces. The fixed 4 GB never moves — buildup can't touch it. This Is why it's no help when the weights alone don't fit. But the activation term, the part that broke the budget, is now a sixteenth of its size. The cost is time: sixteen sequential passes take about if sixteen passes. It buys capability, not throughput, and the one thing it does *not* copy is batch-norm statistics, which are found over the micro-batch and so behave like the small batch but many steps you build up.
Two setup details cause most buildup bugs, and both fail quietly. The loss must be cut by the number of buildup steps before the backward pass, or the built up grad is a sum where the tuner expects a mean and the effective learning rate is silently times'd by the step count. And the grads must be zeroed after the tuner step, not after each backward pass — zero them each pass and the technique is defeated with no error at all, training proceeding with just the small batch it was meant to escape. Activation checkpointing and mixed precision attack the same activation memory from other angles and combine with it. Model parallelism attacks the fixed part that buildup can't.
Where it goes wrong with Gradient Accumulation
Losing to divide the loss. Silently times the learning rate.
Zeroing grads each pass. Silently defeats the technique.
Using it with batch norm at a tiny micro-batch. Real quality loss.
Expecting a speed-up. It buys capability, not throughput.
Questions people ask about Gradient Accumulation
What does gradient accumulation do? It sums gradients over several small batches and applies one optimiser step. You get the effect of a large batch without the memory.
Is it identical to a real large batch? For the gradient, yes, once you divide the loss by the number of accumulation steps. Batch normalisation statistics still differ, since they see only the small batch.
Why must I scale the loss? Gradients add up across micro-batches. Divide each loss by the step count, or the update will be that many times too big.
When do I call optimizer.step()? Only every N micro-batches, followed by zero_grad. Calling backward each time without zeroing is what accumulates.
Does it make training slower? Each update takes N forward and backward passes, so wall time per update grows. Throughput per sample is about the same.
How does it interact with the learning rate? Treat the effective batch size as micro-batch times steps and scale the learning rate as you would for that larger batch.
A large batch on a small GPU
Building up grads over a few small batches before stepping is arithmetically the same to one large batch. This proves the identity, then shows the two things that aren't the same.
example_01.pyNumPy
import numpy as np
rng = np.random.default_rng(0)
N, D = 64, 5
X = rng.normal(size=(N, D))
w_true = rng.normal(size=D)
y = X @ w_true + rng.normal(0, 0.3, N)
w = np.zeros(D)
def grad(Xb, yb, wv):
return 2 * Xb.T @ (Xb @ wv - yb) / len(Xb)
print("64 rows. compute the gradient two ways.")
print()
big = grad(X, y, w)
print("one batch of 64:")
print(" ", np.round(big, 8))
print()
acc = np.zeros(D)
STEPS = 4
for i in range(STEPS):
chunk = slice(i * 16, (i + 1) * 16)
acc += grad(X[chunk], y[chunk], w) / STEPS
print("four batches of 16, each gradient divided by 4 and summed:")
print(" ", np.round(acc, 8))
print(" max difference: %.2e" % np.abs(big - acc).max())
print()
print("identical. that is the whole technique: the mean over 64 rows is the")
print("mean of four means over 16 rows, as long as the batches are equal")
print("sized and you divide by the number of accumulation steps.")
print()
print("the division is the part people get wrong. without it:")
wrong = sum(grad(X[i * 16:(i + 1) * 16], y[i * 16:(i + 1) * 16], w)
for i in range(STEPS))
print(" summed without dividing:", np.round(wrong[:3], 4), "...")
print(" correct: ", np.round(big[:3], 4), "...")
print(" ratio: %.1f -- you have silently multiplied your learning rate by %d."
% (np.abs(wrong).mean() / np.abs(big).mean(), STEPS))
Output
The next block shows what accumulation buys in memory, and the two things it does not preserve.
example_02.pyNumPy
import numpy as np
rng = np.random.default_rng(0)
X = rng.normal(size=(64, 5))
print("what accumulation buys, in memory terms. activation memory scales with")
print("batch size; gradient memory does not:")
print("%14s %18s %18s" % ("micro-batch", "activations held", "effective batch"))
for micro in (64, 16, 8, 4, 1):
print("%14d %18s %18d" % (micro, "%dx" % micro, 64))
print(" the gradient buffer is the same size in every row -- one number per")
print(" parameter. only the activations kept for the backward pass shrink.")
print()
print("and now the two things that are NOT identical.")
print()
print("1. BatchNorm. its statistics come from the micro-batch, not the")
print(" accumulated one, so accumulation does not restore large-batch")
print(" behaviour there:")
for micro in (64, 16, 4):
b = X[:micro]
print(" micro-batch of %2d: feature-0 mean %+.4f, sd %.4f"
% (micro, b[:, 0].mean(), b[:, 0].std()))
print(" those are the numbers BatchNorm would use. they are noisier the")
print(" smaller the micro-batch, and accumulation cannot fix it.")
print(" LayerNorm and GroupNorm are unaffected, which is one more reason")
print(" transformers are easier to train this way.")
print()
print("2. speed. four passes of 16 do the same arithmetic as one pass of 64,")
print(" but with four times the kernel launches and less parallelism. you")
print(" are trading wall-clock time for memory, not getting both.")
Output
Finally, the standard accumulation loop in pseudocode.
example_03.pyNumPy
ACCUM = 4
print("the loop, written out, because the off-by-one is easy:")
print(" optimizer.zero_grad()")
print(" for i, batch in enumerate(loader):")
print(" loss = criterion(model(batch.x), batch.y) / ACCUM")
print(" loss.backward() # gradients ADD by default")
print(" if (i + 1) % ACCUM == 0:")
print(" optimizer.step()")
print(" optimizer.zero_grad()")
print()
print("three things to check: divide the loss by ACCUM, do not zero the")
print("gradients between micro-batches, and handle the final partial group")
print("-- if the loader ends mid-accumulation those gradients are discarded.")
print()
print("and if you use a learning rate schedule, step it once per OPTIMIZER")
print("step, not once per micro-batch, or your schedule runs %d times too fast."
% ACCUM)
Output
Check yourself
0 of 3
Answer without scrolling back up.
Why is the accumulated update identical to a large-batch update?
Identical up to floating-point ordering, not merely similar. That linearity is the whole justification.
Which part of training memory does accumulation reduce?
Activations scale with batch size; weights, gradients and optimiser state do not. When the weights alone do not fit, accumulation is not the answer.
Which layer type does gradient accumulation genuinely change?
Its statistics come from whatever is in the forward pass - the micro-batch. Layer norm normalises per example and is indifferent to the batch, which is one reason transformers use it.
Cheat sheet
Gradient Accumulation
Accumulate gradients over several small forward passes before one optimiser step and get large-batch training on a device that cannot hold one.
DEEP LEARNING · vizlearn.in/deep_learning/gradient_accumulation.html
Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.