Gradient Accumulation

The problem

Language models are trained with batches far too large to fit in memory at once. Gradient accumulation splits a batch into smaller micro-batches: each gets its own forward and backward pass, their gradients are combined, and the optimizer takes one step. Done right, the result is exactly the gradient of the whole batch.

You are given grad_fn(x, y), which returns (loss, grad) for one micro-batch: the mean loss over its targets, and the gradient of that mean (a NumPy array).

Combining means correctly

The loss of the whole batch is the mean over all its targets. If micro-batch $i$ has $n_i$ targets and mean loss $\text{mean}_i$, its targets add up to $n_i \cdot \text{mean}_i$, so:

where:

• $n_i$ is the number of targets of micro-batch $i$ (y.size);
• $n = \sum_i n_i$ is the total number of targets;
• $n_i / n$ is the micro-batch's share of the targets: its weight.

The gradient is a sum of the same kind, so the same weights apply to the micro-batch gradients.

Example. Two micro-batches: the first has 3 targets and mean loss $2.0$, the second 1 target and mean loss $4.0$. So $n = 4$.

• Right: weights $3/4$ and $1/4$, loss $0.75 \cdot 2.0 + 0.25 \cdot 4.0 = 2.5$. Check: the four targets have losses adding up to $3 \cdot 2.0 + 4.0 = 10$, and $10 / 4 = 2.5$.
• Wrong: the plain average of the two means, $(2.0 + 4.0) / 2 = 3.0$, counts the single target of the second micro-batch as much as the three of the first.

When all micro-batches have the same size, the weights are equal and the plain average is right. The last micro-batch is often smaller, and that is where the bug hides.

What to implement

accumulate(grad_fn, micro_batches): micro_batches is a list of (x, y) pairs. Return (loss, grad) for the whole batch: the mean loss over all targets and its gradient.

Requirements

• Weight each micro-batch by y.size divided by the total, as above.
• Call grad_fn exactly once per micro-batch.

Hint

Compute the total number of targets first, in a separate pass over the list: then each micro-batch's weight is known when you call grad_fn on it.