batch size

Batch size is the number of training examples averaged into a single gradient update, and at the scale of modern language models it stops being a small modeling choice and becomes a distributed-systems problem: OpenAI trained GPT-3 with a batch of roughly 3.2 million tokens per step, and Google's PaLM paper (2022) describes deliberately increasing the batch over the course of training, from roughly 1 million tokens up through 2 million and then 4 million as the run progressed. No single GPU holds a batch that size in memory, so the number a paper reports is actually assembled from many smaller pieces spread across hardware and time, as the pipeline below shows.

Each GPU gets a micro-batch
e.g. 4 sequences — as many as memory allows
Local forward + backward pass
gradients computed independently per GPU
All-reduce averages gradients
every GPU ends up with the same averaged gradient
Optional: accumulate first
repeat micro-batches k times before syncing, to stretch memory further
Sync every step
no accumulation — simplest, highest communication cost
Optimizer step
one weight update per effective batch, not per micro-batch

The effective batch size is a product of three numbers, and only one of them is a modeling choice. Effective batch size equals the per-GPU micro-batch size, multiplied by the number of GPUs, multiplied by the number of gradient-accumulation steps — so a team running 64 GPUs with a micro-batch of 4 and 8 accumulation steps reaches an effective batch of 2,048 sequences without ever holding more than 4 sequences in memory on any single device at once. Gradient accumulation exists purely to buy effective batch size when GPU memory, not the desired batch, is the binding constraint: it runs several micro-batch forward-backward passes, sums the gradients instead of applying them immediately, and only calls the optimizer once the target count is reached.

Three Ways to Reach the Same Effective Batch effective batch = per-GPU micro-batch × GPU count × accumulation steps — target: 2,048 sequences More GPUs 512 GPUs micro-batch 4 1 accumulation step 512 × 4 × 1 = 2,048 Fastest wall-clock time. Highest hardware cost. 512-way all-reduce every single step Gradient accumulation 64 GPUs micro-batch 4 8 accumulation steps 64 × 4 × 8 = 2,048 8x slower wall-clock than the 512-GPU run. 64-way all-reduce, 8x less often Bigger micro-batch 64 GPUs micro-batch 32 1 accumulation step 64 × 32 × 1 = 2,048 Needs 8x the activation memory per GPU. Often impossible — GPU runs out of memory Same effective batch, same gradient estimate — the choice is purely about hardware, not the model.

Doubling the batch size does not automatically mean the same number of training examples teaches the model twice as fast. Below a task-and-model-specific threshold — the critical batch size — a larger batch gives a less noisy gradient estimate and genuinely supports a proportionally larger learning rate and fewer total steps; past that threshold, extra examples in the batch mostly cancel out noise that was already small, so step count stops falling and the only effect of a bigger batch is more compute spent per step for roughly the same number of steps saved. McCandlish and colleagues' 2018 empirical study of this exact trade-off is why "just use the biggest batch your hardware allows" is treated as a default to test, not a rule to trust blindly.

Large-batch training has a documented generalization cost, not just a speed benefit. Keskar and colleagues' 2016 paper "On Large-Batch Training for Deep Learning" showed that models trained with large batches (8,192 examples) on CIFAR-10 and ImageNet converged to sharper minima and generalized measurably worse than the same architecture trained with small batches (256 examples), even when both reached similar training loss — the smaller batch's gradient noise appears to act as an implicit regularizer that helps the optimizer avoid narrow, poorly-generalizing solutions. This is why practitioners who scale batch size up for throughput also scale the learning rate up to compensate (the linear scaling rule) and add learning-rate warmup, rather than treating a bigger batch as a pure win.

Batch size sits at the intersection of three different constraints that rarely point the same direction: statistics wants it small enough that gradient noise still helps generalization, optimization wants it large enough to approach the critical batch size where steps are cheapest, and hardware wants it as large as fits in memory to amortize communication overhead — which is exactly why gradient accumulation exists, as the one lever that lets a team hit the statistically and optimally chosen batch size without that choice being overridden by whatever a single GPU's memory happens to allow.

Go deeper with CFSGPT

Get AI-powered deep-dives, save terms, and run advanced simulations — free account.

Create Free Account