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.