Gradient Accumulation

Take one optimiser step on a batch that is too large to run at once.

A GPU often cannot hold a full batch. The usual fix is gradient accumulation: split the batch into micro-batches, run backward on each one, and let the gradients add up in .grad. Then take one optimiser step. Done right, the step is exactly the same as one step on the full batch.

Write accumulate_step(model, optimizer, loss_fn, X, y, micro_batch):

1. Clear the old gradients.
2. Split X and y into consecutive chunks of micro_batch rows. The last chunk can be smaller.
3. For each chunk, compute the loss and call backward().
4. Call optimizer.step() once.
5. Return the full-batch loss as a Python float.

The key detail is the scale. loss_fn returns the mean over its chunk. The full-batch loss is the mean over all n rows, so each chunk's loss must be weighted by len(chunk) / n before backward():

full_loss = sum over chunks of  loss_fn(chunk) * len(chunk) / n

Dividing by the number of chunks instead is only correct when all chunks are the same size.

Function signature

def accumulate_step(model, optimizer, loss_fn, X, y, micro_batch):

Example

Input

X has 10 rows, micro_batch = 4

Output

the model runs on chunks of 4, 4 and 2 rows,
the parameters end where one full-batch step puts them,
and the return value is the full-batch loss before the step

Reasoning

The chunk losses are weighted by 4/10, 4/10 and 2/10. These weights turn three chunk means into one mean over all 10 rows, so the summed gradient is the full-batch gradient.

Constraints and notes

• Call the model once for each chunk. The tests count the calls with a forward hook, so do not run the full batch as well.
• .grad may hold old values when your function is called.
• loss_fn can be any loss with mean reduction, for example nn.MSELoss() or nn.CrossEntropyLoss().
• If micro_batch is at least n, there is one chunk.