Batching, Gradient Accumulation, and Throughput

How batch size, accumulation, optimizer steps, and hardware utilization interact during training.

Machine Learning · Training

Batch size is both a statistical choice and a systems choice. It changes gradient noise, memory use, device occupancy, synchronization cost, and the meaning of an optimizer step.

Effective batch size

For data-parallel training:

\[B_{eff} = B_{device} \times N_{devices} \times K_{accum}\]

where $B_{device}$ is the per-device microbatch, $N_{devices}$ is the number of workers, and $K_{accum}$ is the number of accumulation steps before an optimizer update.

If each microbatch loss is averaged, accumulate gradients as:

\[g = \frac{1}{K_{accum}}\sum_{k=1}^{K_{accum}} g_k\]

If the training loop calls loss.backward() on already averaged microbatch losses and never divides by $K_{accum}$, the update is scaled up by $K_{accum}$. That may look like a learning-rate change, not an accumulation bug.

For token losses, normalize by total non-padding tokens across the effective batch, not by the number of microbatches. Otherwise batches with different padding patterns get different weight.

Accumulation is not always equivalent

Gradient accumulation can match a larger batch for plain SGD when these conditions hold:

the same examples are used,
gradients are averaged before the optimizer step,
optimizer state updates once per effective batch,
batch-dependent layers see equivalent statistics,
stochastic operations use equivalent randomness assumptions.

The equivalence breaks when BatchNorm uses microbatch statistics. It can also break when the learning-rate scheduler advances per microstep instead of per optimizer step, when gradient clipping is applied before accumulation, or when optimizer moments update more often than intended.

For Adam-like optimizers, the update uses moment estimates:

\[m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t\] \[v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2\]

If $t$ advances per microbatch instead of per effective batch, bias correction and scheduler timing change.

Throughput limits

Small batches often underuse accelerators because kernel launch overhead, memory latency, and host input pipelines dominate. Larger batches improve arithmetic intensity until memory, communication, or generalization becomes the limiter.

For synchronous data parallelism, each step includes compute plus gradient synchronization. A rough step model is:

\[T_{step} \approx T_{compute}(B_{device}) + T_{allreduce}(|\theta|)\]

Accumulation increases compute per synchronization, which can improve throughput when all-reduce is expensive. It does not reduce activation memory for a single microbatch, but it avoids storing activations for the whole effective batch at once.

Learning-rate scaling

Large-minibatch SGD often uses a linear scaling rule as a starting point:

\[\eta_{new} = \eta_{base}\frac{B_{new}}{B_{base}}\]

This rule is empirical, not a law. Warmup is commonly used because early training can be unstable when a large learning rate is applied before representations settle.

Practical checks

Log optimizer steps, not just dataloader iterations.
Verify scheduler stepping: per microbatch or per optimizer step.
Divide accumulated gradients or scale the microbatch loss intentionally.
Apply clipping after accumulation if the intended bound is on the effective-batch gradient.
In AMP, unscale gradients only after all accumulation for that optimizer is complete.
Track examples/sec and tokens/sec separately from loss quality.

Design guidance