Large-Batch Training and Learning-Rate Scaling

"They gave me a thousand friends to share the work, then doubled my step size to match. For a while I sprinted. Then I tripped over my own enthusiasm and they invented warmup."

An Optimizer That Scaled Its Learning Rate

The previous section reduced the bytes moved per synchronization step. This one asks a complementary question: as you add workers to a synchronous data-parallel job, what actually changes for the optimizer, and how do you keep it converging? The answer is entirely about batch size. With $K$ workers each computing a gradient over a local batch of $B_{\text{loc}}$ examples and an all-reduce averaging them, the optimizer sees one update built from a global batch of $B = K \, B_{\text{loc}}$ examples. Doubling the worker count doubles $B$. Everything stable and everything fragile about large-scale synchronous training follows from that one identity, so we start there and build outward to the rules that tame it.

1. More Workers Means a Bigger Batch, Which Is Weak Scaling Beginner

Hold the per-worker batch fixed and add workers: the global batch grows linearly, and the number of optimizer updates per pass over the data shrinks by the same factor. This is the defining signature of weak scaling, where the total problem size grows with the worker count rather than staying fixed, the regime introduced in Section 3.3. Strong scaling would keep the global batch constant and shrink each worker's share, which quickly starves the accelerators; the dominant practice in distributed training is instead to keep each accelerator well-fed at a fixed local batch and let the global batch expand. So the question "how many workers?" is, for the optimizer, identical to the question "how large a batch can I train with and still converge well?"

That equivalence is why this topic belongs in a chapter on optimization rather than systems. The communication machinery (the all-reduce of Chapter 4) is indifferent to batch size; it averages whatever gradients it is given. The optimizer is not indifferent. A gradient averaged over $B$ examples is an estimate of the true gradient whose variance falls roughly like $1/B$. Larger batches give cleaner gradients, and a cleaner gradient is one you can trust to take a larger step along. The rest of this section turns that intuition into a rule, then shows where the rule breaks.

2. The Linear Scaling Rule and Why Warmup Is Mandatory Intermediate

The practical recipe that made large-batch training routine is the linear scaling rule of Goyal et al. (2017): when you multiply the batch size by a factor, multiply the learning rate by the same factor. Concretely, if a reference batch $B_0$ trains well at learning rate $\eta_0$, then a batch of size $B$ should use

The heuristic behind it is short. Take $K$ consecutive small-batch SGD steps with learning rate $\eta_0$ on batches $\mathcal{B}_1, \dots, \mathcal{B}_K$. If the parameters barely move across those $K$ steps, the gradients are evaluated at nearly the same point, and the net update is approximately $-\eta_0 \sum_{j=1}^{K} \nabla \ell_{\mathcal{B}_j}(w)$. One large-batch step on the union $\mathcal{B}_1 \cup \dots \cup \mathcal{B}_K$ with learning rate $\eta$ produces $-\eta \cdot \frac{1}{K}\sum_{j=1}^{K}\nabla \ell_{\mathcal{B}_j}(w)$. Matching the two requires $\eta = K \eta_0$, exactly the linear rule. The rule is an approximation, and it names its own failure mode: it holds only while the assumption "the parameters barely move" holds.

Early in training the parameters move a great deal, the loss surface is steep and curved, and the linear-rule learning rate (now $K$ times larger) overshoots and can diverge on the very first steps. The fix is a warmup: start the learning rate small and ramp it linearly to the scaled value over the first few hundred updates, by which point the parameters have settled into a region where the matching argument is valid. Warmup is not a cosmetic safety margin; it is the repair for the one place the derivation is known to be invalid. Skip it at large $K$ and the run frequently diverges before it begins.

3. The Critical Batch Size: Where Adding Workers Stops Helping Intermediate

The linear scaling rule cannot hold forever, because no fixed number of training examples needs an unboundedly large step. Below some threshold, doubling the batch (and the learning rate) roughly halves the number of optimizer steps to a target loss, so you reach the target in about the same number of passes over the data and the extra workers translate into real speedup. Above that threshold, doubling the batch no longer halves the steps; the gradient is already clean enough that lowering its variance further does little, and you are in the regime of diminishing returns. McCandlish et al. (2018) named this threshold the critical batch size and gave it a predictor: it scales with the gradient noise scale, a measurable ratio of gradient variance to gradient magnitude that tends to grow as training proceeds.

The critical batch size is the optimization counterpart of the serial bottleneck in Section 3.5: just as Amdahl's law caps speedup once the parallelizable fraction is exhausted, the critical batch size caps useful worker count once the gradient-variance reduction is exhausted. Past it, adding workers grows the batch beyond what the optimizer can convert into faster convergence, so wall-clock-per-epoch keeps falling but epochs-to-target stops falling, and total time-to-train flattens or worsens. Knowing roughly where this knee sits for your model is what separates a worker count that pays for itself from one that burns budget. Figure 10.8.1 sketches the shape.

4. A From-Scratch Demonstration of the Knee Intermediate

The claims above are testable in pure Python without any cluster. The code below trains a logistic-regression model with plain mini-batch SGD on a fixed synthetic dataset, sweeping the global batch from 32 to 4096. For each batch it measures epochs-to-target-loss twice: once with a fixed learning rate (the naive baseline) and once with the linear scaling rule plus warmup. Counting epochs (passes over the data) is the right cost unit, because a larger batch means fewer optimizer updates per epoch, so a fixed learning rate makes less progress per pass while the scaled rule compensates.

import numpy as np

rng = np.random.default_rng(7)
N, d = 20_000, 30
X = rng.standard_normal((N, d))
w_star = rng.standard_normal(d)
p = 1.0 / (1.0 + np.exp(-(X @ w_star)))
y = (rng.random(N) < p).astype(np.float64)        # noisy binary labels

def loss_grad(w, Xb, yb):
    pr = 1.0 / (1.0 + np.exp(-(Xb @ w)))
    g = Xb.T @ (pr - yb) / len(yb)
    L = -np.mean(yb*np.log(pr+1e-12) + (1-yb)*np.log(1-pr+1e-12))
    return L, g

full_loss = lambda w: loss_grad(w, X, y)[0]

# Demanding target a little above the optimum, so reaching it takes many epochs.
w_opt = np.zeros(d)
for _ in range(4000):
    w_opt -= loss_grad(w_opt, X, y)[1]
TARGET = full_loss(w_opt) + 0.010

base_bs, base_lr, MAX_EPOCHS = 32, 0.05, 80          # reference batch and its tuned lr

def epochs_to_target(B, scale_lr):
    lr = base_lr * (B / base_bs) if scale_lr else base_lr   # linear scaling rule
    warmup = 5 * max(1, B // base_bs) if scale_lr else 0    # warmup grows with scaling
    w, steps = np.zeros(d), 0
    for epoch in range(1, MAX_EPOCHS + 1):
        perm = rng.permutation(N)
        for s in range(0, N, B):
            _, g = loss_grad(w, X[perm[s:s+B]], y[perm[s:s+B]])
            cur = lr * (steps+1)/warmup if warmup and steps < warmup else lr
            if not np.all(np.isfinite(g)): return None       # diverged
            w -= cur * g
            steps += 1
        if full_loss(w) <= TARGET:
            return epoch
    return None                                              # never reached target

for B in [32, 64, 128, 256, 512, 1024, 2048, 4096]:
    ef = epochs_to_target(B, scale_lr=False)
    es = epochs_to_target(B, scale_lr=True)
    fmt = lambda e: ">80 (stalled)" if e is None else str(e)
    print(f"{B:>6} | fixed-lr {fmt(ef):>13} | scaled-lr {fmt(es):>13}")

    32 | fixed-lr             2 | scaled-lr             2
    64 | fixed-lr             4 | scaled-lr             2
   128 | fixed-lr             8 | scaled-lr             3
   256 | fixed-lr            16 | scaled-lr             3
   512 | fixed-lr            31 | scaled-lr             4
  1024 | fixed-lr            61 | scaled-lr             6
  2048 | fixed-lr >80 (stalled) | scaled-lr            11
  4096 | fixed-lr >80 (stalled) | scaled-lr            22

Read the two columns together and the story of this section is in the numbers. The naive column shows why you cannot simply pour workers into a synchronous job and hope: without scaling the learning rate, a bigger batch is strictly worse. The scaled column shows the linear rule doing its job, holding epochs-to-target nearly constant while the batch grows eightfold, then the same column reveals the knee, where even correct scaling cannot recover the $1/B$ ideal because the gradient is already low-variance. The flattening past batch 512 is the demonstration, in miniature, that the useful worker count is bounded by the optimization problem, not by the cluster you can rent.

5. Layer-Wise Adaptive Methods: Pushing the Knee Outward Advanced

The critical batch size is not a hard constant; it depends on the optimizer. The linear scaling rule applies one global learning rate to every parameter, but in deep networks the ratio of update magnitude to weight magnitude varies enormously across layers, and at very large batch a single global rate is simultaneously too large for some layers and too small for others, which is what makes naive scaling diverge well before the statistical knee. Layer-wise adaptive methods attack this directly by giving each layer its own effective step. LARS (You, Gitman, Ginsburg, 2017) rescales each layer's update so its norm is a fixed fraction of that layer's weight norm, which let ResNet-style networks train at batches in the tens of thousands. LAMB (You et al., 2019) carries the same trust-ratio idea into the Adam family and is the method that trained BERT to its target in 76 minutes at a batch of 32k, a result that would diverge instantly under a single global linearly-scaled rate.

The lesson for distributed training is that the optimizer choice sets how far you can scale out before the knee bites. Plain momentum SGD with linear scaling and warmup is robust to moderate batches; LARS and LAMB extend the usable range by an order of magnitude or more for the architectures they were designed for, by replacing one fragile global step with many self-calibrated per-layer steps. They do not abolish the critical batch size, they relocate it, which is exactly the leverage you want when the alternative is leaving expensive accelerators underused.

# Run with: torchrun --nproc_per_node=8 train.py
import torch, torch.distributed as dist
from torch.optim.lr_scheduler import LinearLR

dist.init_process_group("nccl")
K = dist.get_world_size()                         # number of workers == batch multiplier

base_lr, base_bs = 0.05, 32
lr = base_lr * (K * local_bs) / base_bs           # linear scaling rule from the world size
opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)

# Linear warmup over the first 500 steps, then hold the scaled rate.
warmup = LinearLR(opt, start_factor=1/K, total_iters=500)
# model = DistributedDataParallel(model)  # averages gradients -> global batch K*local_bs

6. The Practical Ceiling and What Comes Next Intermediate

Putting the pieces together gives a sober rule for sizing a synchronous data-parallel job. Add workers while the global batch stays below the critical size, scaling the learning rate linearly and warming it up, and use a layer-wise adaptive optimizer if a single global rate begins to diverge before the statistical knee. Stop adding workers near the knee, because past it the extra hardware reduces per-epoch wall-clock without reducing epochs-to-target, so total training time stops improving and cost rises with no return, an Amdahl-shaped ceiling on the optimization side to match the communication-side ceiling of Section 3.5. The useful worker count is the smaller of what your network can sustain and what your optimizer can absorb, and Output 10.8.1 shows the second limit can bind well before the first.

We have treated communication as something to minimize and batch size as something to bound, but we have not yet asked how little communication is fundamentally required to optimize across machines, regardless of cleverness. That lower-bound question, how many bits and rounds any distributed optimizer must pay, is the subject of the next section. It turns the engineering trade-offs of this chapter into the theory that says which of them are avoidable and which are laws.