Gradient accumulation: simulating a larger batch size than memory allows

The problem, from first principles

[03-language-model-assembly/05-training-loop]'s train_output_head_one_step computes a gradient from ONE batch and immediately applies it. Larger batches generally give a more stable, less noisy gradient estimate, but a batch's memory footprint grows with its size, and eventually a desired batch size simply doesn't fit in available memory all at once, a real, common constraint when training large models. Gradient accumulation resolves this without needing more memory: split the desired large batch into several smaller MICRO-batches that DO fit individually, compute each micro-batch's gradient SEPARATELY (never holding more than one micro-batch in memory at a time), average those gradients together, and apply only ONE weight update using the averaged result, mathematically equivalent to having computed the gradient on the full large batch directly.

The correctness of this trick rests on one precise mathematical fact: for a MEAN-reduced loss ([03-next-token-cross-entropy]'s default), the gradient of the mean over a large batch equals the MEAN of the gradients of the mean over equal-sized smaller groups that partition it (a nested average of equal-sized groups is exactly the overall average). Get the averaging wrong (accumulate via SUM instead of MEAN, say) and the effective step size silently changes with the number of micro-batches, a subtle, easy-to-introduce training bug.

From theory to code

Implement compute_output_head_gradient(hidden_states, token_ids, output_weight) ([05-training-loop]'s gradient computation, without the update), accumulate_gradients(gradients) (the elementwise mean across micro-batch gradients), and train_with_gradient_accumulation(micro_batches, output_weight, lr), combining them into one full accumulation-then-update cycle.

Constraints

• Every micro-batch's gradient is computed INDEPENDENTLY, with output_weight held FIXED across all of them (no intermediate updates between micro-batches).
• accumulate_gradients averages (MEAN), never sums, the micro-batch gradients.
• Exactly ONE weight update happens, using the accumulated (averaged) gradient, after ALL micro-batches have been processed.
• For equal-sized micro-batches, the resulting update must be IDENTICAL (up to floating-point precision) to computing the gradient directly on the full concatenated batch in one shot.

Hints

gradients = [compute_output_head_gradient(h, t, output_weight)[0] for h, t in micro_batches]
accumulated_grad = np.mean(gradients, axis=0)
updated_weight = output_weight - lr * accumulated_grad