gradaccumbias

the question

Does naive gradient accumulation over a mean-reduction loss compute the true large-batch gradient when micro-batches have unequal token counts?

what came out

No. It computes a token-count reweighted gradient biased by exactly sum_i g_i(1/(k n_i)-1/N), vanishing only for equal counts. At counts [3,7,2,11] it differs by 90% of magnitude, and a lone token trains at 50.5x rate. Summing token gradients and dividing once by N recovers the true gradient bit for bit.

→ read the code and re-run it