Training at Scale

gradient accumulation

Large-batch training is often desirable for stability and throughput, but the batch you want may not fit in memory at once. Gradient accumulation decouples the statistical batch from the physical one. You split the desired effective batch into several micro-batches, run a forward and backward on each, and add the resulting gradients into a buffer without taking an optimizer step. Only after the last micro-batch — when the accumulated gradient equals what one big batch would have produced — do you apply a single update.

Because each micro-batch is processed and freed before the next, peak activation memory is set by the micro-batch size, not the effective batch size, so you can reach an arbitrarily large effective batch on fixed hardware at the cost of proportionally more sequential steps. The subtlety is correctness of averaging: the loss must be scaled (or gradients divided by the accumulation count) so the accumulated gradient is the mean, not the sum; and in data-parallel training you should suppress the all-reduce on every micro-batch except the last to avoid averaging partial gradients prematurely.

Gradient accumulation is also the mechanism that fills a pipeline: the micro-batches it creates are exactly the units that flow through pipeline stages to shrink the bubble. It is nearly free in memory but not in time — more accumulation steps mean more forward/backward passes per update, so it raises effective batch size without raising hardware utilization.

g = \frac{1}{K}\sum_{k=1}^{K} \nabla_\theta \mathcal{L}(B_k),\quad \theta \leftarrow \theta - \eta\, g

Averaging K micro-batch gradients before a single step reproduces the gradient of one batch of size |B_k|·K.

Gradient accumulation is only equivalent to a true large batch up to batch-statistics layers: components like batch normalization compute statistics per micro-batch and so do not reproduce the large-batch behavior, though transformers' layer norm is unaffected.

Also called
micro-batching梯度累積