Batch Size Math for LLM Training

What batch size actually counts

Three numbers get called ‘batch size,’ so fix the vocabulary. The microbatch is what one device forwards at once — a memory constraint. The global batch is what one optimizer step consumes:

B_global = microbatch × grad_accum_steps × data_parallel_workers
tokens_per_step = B_global × seq_len

Only B_global appears in the optimization math; the split into microbatches and accumulation steps is memory engineering that leaves the update mathematically identical (the sibling article on gradient accumulation covers the mechanics and the BatchNorm caveat). For language models the honest unit is tokens per step, not sequences — 512 tokens and 8192 tokens carry very different gradient signal. Frontier runs land in the millions: GPT-3 trained at about 3.2M tokens per step, Llama-family runs at roughly 4M, while a 100M-parameter model on a laptop might use 64k. That thousand-fold spread tracks a real quantity, which is what the rest of this article derives.

The gradient is an estimator, and its variance falls as 1/B

The loss you actually want to descend is the expectation over the data distribution, L(θ) = E_x[ℓ(x, θ)], with true gradient G = ∇L. You never see G. You see a batch average of per-example gradients:

g_B = (1/B) Σ_{i=1..B} g_i,    E[g_B] = G
Cov(g_B) = Σ / B,    where Σ = Cov(g_i)

The estimator is unbiased at every batch size; what changes is the noise. The expected squared error is E[|g_B − G|^2] = tr(Σ) / B; divide by the squared signal |G|^2 for a dimensionless ratio:

noise / signal = tr(Σ) / (B · |G|^2) = B_simple / B
B_simple := tr(Σ) / |G|^2

That ratio is the whole story. Doubling B halves the gradient’s variance but costs twice the compute, and B_simple / B tells you whether you are buying signal or paying for decimal places you cannot use.

The gradient noise scale

B_simple is the crude version. The sharper quantity, from the empirical large-batch model of McCandlish et al. (2018), asks how much the loss improves per step. Expand to second order around θ for a step −η g_B and take the expectation:

E[ΔL] = −η |G|^2 + (η^2/2) · ( GᵀHG + tr(HΣ)/B )

The tr(HΣ)/B term is pure noise penalty, and it is the only place B appears. Optimizing over η gives the noise scale B_noise = tr(HΣ) / (GᵀHG) and a strikingly clean result:

ΔL_opt = ΔL_max / (1 + B_noise / B)
η_opt  = η_max  / (1 + B_noise / B)

Read the limits. For B « B_noise, improvement per step is roughly proportional to B — every extra example buys real progress. For B » B_noise, improvement saturates at ΔL_max and extra examples buy nothing. B_noise is the elbow. In practice B_simple stands in for it, needing no Hessian. Both are estimable mid-run from the gap between per-worker and all-reduced gradient norms, since E[|g_B|^2] = |G|^2 + tr(Σ)/B at two batch sizes gives two equations in two unknowns.

Critical batch size: the steps-versus-examples hyperbola

Now turn per-step improvement into a training budget. Let S be the optimizer steps needed to reach a target loss and E = B · S the tokens consumed. Substituting the expression above gives two equations and one tidy identity:

S / S_min = 1 + B_crit / B
E / E_min = 1 + B / B_crit           (E_min = S_min · B_crit)

(S/S_min − 1) · (E/E_min − 1) = 1

The Pareto frontier of ‘serial time’ against ‘total compute’ is a hyperbola, and B_crit is the point where you pay 2× on both axes. S_min is the irreducible step count at infinite batch, E_min the irreducible token count at infinitesimal batch; you cannot beat either. Nothing here says a bigger batch is worse — it says a bigger batch converts tokens into fewer sequential steps, at a rate that collapses past B_crit. If your bottleneck is wall-clock and you have spare accelerators, overshooting is a rational purchase; if it is data or FLOPs, it is waste.

A worked example

Take a run with B_crit = 1M tokens per step, S_min = 50,000 steps, so E_min = 5×10^10 tokens. The table falls straight out of the two formulas:

Going 250k → 1M cuts steps by 2.5× for 1.6× the tokens: an excellent trade if you have the hardware. Going 4M → 16M cuts steps by 18% for 3.4× the tokens: almost pure waste. The knee at B_crit is sharp, which is why production runs cluster within a small factor of their estimated critical batch rather than using every GPU-hour of parallelism available.

Linear scaling, and why warmup is mandatory

Change the batch and you must change the learning rate, or you have changed the optimizer. The classic argument (Goyal et al., 2017): take k SGD steps with batch B and rate η, assuming the gradient barely moves across them:

θ − η Σ_{j=1..k} g_{B,j}  ≈  θ − (kη) · g_{kB}

One step at batch kB with rate kη matches k steps at batch B with rate η. Hence the linear scaling rule, η ∝ B — which is exactly the small-batch limit of the noise-scale formula, since η_opt ≈ η_max · B/B_noise when B « B_noise. The noise scale also says where it stops: past B_crit, η_opt saturates.

The rule’s assumption is worst at initialization, where the surface is sharpest and the gradient moves fastest — which is why large-batch runs diverge in their first few hundred steps. The fix is a gradual warmup: ramp linearly from a small rate to the target over hundreds to thousands of steps. Goyal et al. showed this alone made linear scaling hold to batch 8k; every LLM recipe carries the idea.

Square-root scaling and why Adam is different

Linear scaling is an SGD result, and LLMs are trained with Adam/AdamW. Adam divides by a running estimate of the gradient’s second moment, which already normalizes away part of the magnitude change that linear scaling was compensating for. Per coordinate, the update is roughly η · m / (√v + ε); when noise dominates, √v itself shrinks like 1/√B, so a fixed η already gets part of the boost for free. Keeping the variance of the update constant then gives η ∝ √B — the square-root rule, and empirically the better default for Adam-trained transformers.

Treat both rules as starting points, not laws: they assume the gradient is roughly constant across the merged steps and that you are below B_crit, and neither holds early or at extreme batch. Practical protocol — change B by a factor k, scale η by √k, then sweep ±2×. Note AdamW’s decoupled weight decay does not follow the learning rate, so scaling η silently changes the decay balance unless you compensate.

The large-batch generalization question

Keskar et al. (2017) reported that large-batch training generalizes worse and attributed it to convergence toward sharp minima, where small parameter perturbations cause large loss increases, versus the flat minima that small-batch noise finds. The mechanism is plausible: SGD noise scales as η/B and acts like a temperature, and lowering it lets the optimizer settle into narrower basins.

The follow-up literature qualified this heavily. Most of the gap disappears once you retune the learning rate, add warmup, and hold epochs rather than steps fixed — a fixed-step comparison simply gives the large batch far fewer updates. Adaptive methods (LARS, LAMB) pushed BERT to batch 32k with no quality loss. For LLM pretraining, which is typically single-epoch and undertrained rather than overfit, the binding constraint is the B_crit efficiency curve, not a generalization penalty.

Batch ramps and why the schedule is not flat

If B_crit grows during training, a constant batch size is mis-specified at both ends of the run: too large at the start (wasting tokens on a gradient that was already clean) and too small at the end (wasting steps on a noisy one). The natural response is a batch-size ramp, and this is exactly what large runs do — GPT-3 ramped from roughly 32k tokens per step to 3.2M over the first several billion tokens, and Chinchilla-era runs doubled the batch partway through.

A ramp interacts cleanly with the rest of the recipe if you remember two things. The learning rate must track it — with square-root scaling a 4× batch increase wants roughly a 2× rate increase, on top of your decay schedule. And a ramp changes tokens-per-step, so any schedule defined in steps (warmup, cosine period) silently changes meaning; define schedules in tokens consumed and the ramp becomes a free efficiency win rather than a source of loss spikes.

Hardware pulls one way, a CPU budget pulls the other

Hardware is why the industry drifted to multi-million-token batches. Arithmetic intensity: transformer GEMMs become more compute-bound as the token dimension grows, raising FLOPs utilization. Communication amortization: a gradient all-reduce costs fixed bytes per step, so doubling the batch halves per-token comms. Fixed per-step costs: optimizer updates and kernel launches are paid once per step. The trap is reading throughput as progress — a run at 4× B_crit can post a beautiful tokens-per-second number and still hit the target loss later, because it spent 5× the tokens to save 20% of the steps.

On a CPU the calculus inverts. There is no idle cluster to fill, so there is no reason to exceed B_crit — every token past the knee is wall-clock you never get back. And B_crit is small here: a small model on a modest corpus sits at high loss throughout, so tens of thousands of tokens per step often suffice. Use accumulation to reach a batch that is stable, not one that imitates a frontier recipe three orders of magnitude larger.

Pitfalls that quietly break the math

Several everyday details invalidate the formulas without any error message. Padding: if a nominal batch of 256 sequences is 40% padding, your real batch is 60% of what you think — report tokens and prefer packed sequences. Sequence length: doubling seq_len at fixed sequence count doubles tokens per step, a batch change that needs a learning-rate response. Loss reduction: averaging per-sequence rather than per-token reweights short sequences, changing both gradient and variance. Unscaled rates: changing the number of data-parallel workers changes B_global, so a fixed η is a different experiment. Gradient clipping: a global-norm clip interacts with batch size, since E[|g_B|] falls as B rises — a threshold tuned at small batch may never trigger at large batch, removing the stabilization you relied on.