Gradient accumulation: a big batch on a small GPU

Build a Small GPT From Scratch · Part 4 of 8

Large batches make training steadier, but a free GPU can only hold a small one. Gradient accumulation closes the gap: run several small batches, add up their gradients, and update once.

The idea

The gradient of an average loss over 32 examples is the average of the gradients over smaller groups of them. So instead of pushing 32 examples through at once, you can run four batches of 8, keep adding their gradients together, and only then take one optimizer step. Memory is only needed for 8 examples at a time, yet the update is the one a batch of 32 would have produced.

Checking it on a small problem

This is easy to believe and easy to get wrong, so here it is checked. On a small least-squares problem with 32 examples, the gradient from one big batch and the gradient from four accumulated micro-batches of 8 agree to numerical precision:

import numpy as np
rng = np.random.default_rng(0)
X = rng.normal(size=(32, 5)); y = X @ np.array([1., -2., .5, 0., 3.]) + rng.normal(scale=.1, size=32)
w = np.zeros(5)

def grad(Xb, yb, w):                          # gradient of the mean squared error over a batch
    return 2 * Xb.T @ (Xb @ w - yb) / len(yb)

big = grad(X, y, w)
acc = sum(grad(X[i:i+8], y[i:i+8], w) / 4 for i in range(0, 32, 8))        # divide by the number of micro-batches
bad = sum(grad(X[i:i+8], y[i:i+8], w) for i in range(0, 32, 8))            # forgot the division
print(np.allclose(big, acc), np.allclose(big, bad), np.allclose(bad, 4 * big))
# True False True

The third number shows the classic bug: forget to divide by the number of micro-batches and the gradient comes out exactly four times too large, which behaves like a learning rate four times too high.

In a training loop

In PyTorch the pattern is: scale each micro-batch loss by one over the number of micro-batches, call backward() every time (gradients add up in memory), and call the optimizer step and zero_grad() only after the last micro-batch. The cost of accumulation is time, not memory: you do four forward and backward passes per update instead of one. There is one subtlety: layers whose output depends on the whole batch, such as batch normalisation, do not accumulate this way, but the layer normalisation used in transformers does not depend on the batch and is unaffected.

Open in Colab Run the check yourself and try a different number of micro-batches.

Takeaways

• Accumulating gradients over micro-batches gives the same update as one big batch.
• Divide each loss by the number of micro-batches, or the gradient is too large by that factor.
• It costs extra forward and backward passes, not extra memory.
• It is how a free Colab GPU can imitate the batch sizes of bigger machines.

Next in the series (Part 5): Half the bytes: mixed precision and checkpoints