Learn AI Create the Future — with Pranay
Main site

Multi-GPU and Distributed Training

Scaling the learning rate with the number of GPUs

Adding GPUs multiplies your effective batch size, which makes each gradient less noisy — so the same learning rate now takes steps that are too small, and the usual fix is to multiply it by the world size with a warmup.

Read these first

On this page 5
  1. Why this bites people
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

More GPUs means a bigger batch, a steadier signal, and room to take bigger steps.

Imagine asking people on the street which way the station is. Ask one person and you might get a confident wrong answer, so you take a small careful step. Ask sixteen people and average their answers, and you can stride out with confidence.

Adding GPUs is asking more people at once. Each worker looks at its own slice of data and reports which way to adjust the model. Those reports get averaged.

The average of many reports wobbles less than a single report. A step size chosen for one wobbly report is now too timid.

Why this bites people

Nothing warns you. You move from 1 GPU to 8, the job runs, the loss goes down, and the run looks healthy. The run reaches a given quality more slowly than it should, and nothing inside the run reveals that.

The reverse is worse. Multiply the step size by 8 without care and the first few hundred steps can blow up. The model starts from random weights, where large steps are dangerous.

How it works

1 GPU,  batch 32   ->  one report from 32 examples   -> step size L
8 GPUs, batch 32   ->  average of 8 reports          -> effective batch 256
   each                 = steadier signal
                        -> step size around 8 x L

but for the first few hundred steps:
   start at L and grow smoothly to 8 x L   (the "warmup")

A real example you have seen

Election exit polls. A poll of 200 people swings wildly between samples; a poll of 20,000 barely moves. News channels commit to a confident prediction only once the sample is big enough. Same logic, same reason.

Remember this

  • Adding workers multiplies the effective batch, it does not change the maths.
  • Bigger batch means a steadier gradient, so you can afford a bigger step.
  • The common rule is multiply the learning rate by the world size, with a warmup.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Captured on CPU with the gloo backend. No GPU needed.

Proof that N ranks equal one big batch

effective_batch.py
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="gloo")
rank, world = dist.get_rank(), dist.get_world_size()

torch.manual_seed(0)
full_x, full_y = torch.randn(8, 4), torch.randn(8, 1)    # one batch of 8

# the reference: a single process doing the whole batch of 8
torch.manual_seed(1)
solo = nn.Linear(4, 1)
nn.functional.mse_loss(solo(full_x), full_y).backward()
reference = solo.weight.grad.clone()

# the DDP run: each rank sees 8/world of that same batch
torch.manual_seed(1)
ddp = DDP(nn.Linear(4, 1))
per_rank = 8 // world
mine = slice(rank * per_rank, (rank + 1) * per_rank)
nn.functional.mse_loss(ddp(full_x[mine]), full_y[mine]).backward()

if rank == 0:
    print(f"world_size = {world}, per-rank batch = {per_rank}")
    print("one process, batch 8 :", [round(v, 5) for v in reference.flatten().tolist()])
    print("ddp, batch 8 in total:",
          [round(v, 5) for v in ddp.module.weight.grad.flatten().tolist()])
    print("identical?", torch.allclose(reference, ddp.module.weight.grad, atol=1e-6))

dist.destroy_process_group()
bash
torchrun --nproc_per_node=2 effective_batch.py
torchrun --nproc_per_node=4 effective_batch.py
Output
world_size = 2, per-rank batch = 4
one process, batch 8 : [-0.8099, -0.06596, 0.01549, 0.37768]
ddp, batch 8 in total: [-0.8099, -0.06596, 0.01549, 0.37768]
identical? True

world_size = 4, per-rank batch = 2
one process, batch 8 : [-0.8099, -0.06596, 0.01549, 0.37768]
ddp, batch 8 in total: [-0.8099, -0.06596, 0.01549, 0.37768]
identical? True

DDP with a per-rank batch of b and R ranks produces the identical gradient to one process with a batch of R × b. Not similar — identical, to five decimals, at both world sizes.

So the quantity that governs your learning rate is batch_size × world_size, and batch_size in your config file no longer means what it used to.

Why a bigger batch earns a bigger step

gradient_noise.py
import torch
import torch.nn as nn

torch.manual_seed(0)
X = torch.randn(4096, 6)
y = X @ torch.randn(6, 1) + 0.5 * torch.randn(4096, 1)
model = nn.Linear(6, 1)

def one_gradient(batch):
    idx = torch.randint(0, len(X), (batch,))
    model.zero_grad(set_to_none=True)
    nn.functional.mse_loss(model(X[idx]), y[idx]).backward()
    return model.weight.grad.clone()

print(f"{'batch':>6}  {'spread of the gradient estimate':>32}")
for batch in (8, 32, 128, 512):
    samples = torch.stack([one_gradient(batch) for _ in range(400)])
    print(f"{batch:>6}  {samples.std(dim=0).mean().item():>32.4f}")
Output
 batch   spread of the gradient estimate
     8                            1.6912
    32                            0.8423
   128                            0.4032
   512                            0.2083

Look at the pattern. Each row multiplies the batch by 4 and roughly halves the spread. Noise falls with the square root of the batch, not with the batch itself.

This is the honest tension inside the linear scaling rule. Batch grows 4x, noise falls 2x, and the rule tells you to raise the learning rate 4x anyway. It works in practice over a useful range, and it stops working when the batch gets large enough that noise is no longer what limits your step size.

The rule, and the warmup

scaled_lr.py
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="gloo")
world = dist.get_world_size()

BASE_LR = 0.1                      # the rate you tuned on ONE worker
WARMUP_STEPS = 5

torch.manual_seed(0)
ddp = DDP(nn.Linear(4, 1))
opt = torch.optim.SGD(ddp.parameters(), lr=BASE_LR * world)   # the linear rule

# climb from BASE_LR to BASE_LR * world over the first WARMUP_STEPS
warmup = torch.optim.lr_scheduler.LinearLR(
    opt, start_factor=1.0 / world, end_factor=1.0, total_iters=WARMUP_STEPS)

torch.manual_seed(1)
x, y = torch.randn(4, 4), torch.randn(4, 1)
for step in range(7):
    opt.zero_grad(set_to_none=True)
    nn.functional.mse_loss(ddp(x), y).backward()
    opt.step()
    warmup.step()
    if dist.get_rank() == 0:
        print(f"step {step}: lr = {opt.param_groups[0]['lr']:.4f}")

dist.destroy_process_group()
bash
torchrun --nproc_per_node=4 scaled_lr.py
Output
step 0: lr = 0.1600
step 1: lr = 0.2200
step 2: lr = 0.2800
step 3: lr = 0.3400
step 4: lr = 0.4000
step 5: lr = 0.4000
step 6: lr = 0.4000

The first step is taken at the single-worker value of 0.1. The rate then climbs over five steps and settles at 0.4 — four times the base, for four workers. The printed number is the rate after warmup.step(), which is why the first line already reads 0.16 rather than 0.1.

Real runs use hundreds or thousands of warmup steps, not five. Five keeps the output short enough to read.

Common mistakes

Keeping the same learning rate after adding GPUs. Your run is correct and slower to converge than it should be. This one is invisible without a controlled comparison.

Scaling the rate with no warmup. Early training with random weights takes very large steps and can produce a NaN loss — see debugging NaN loss.

Applying the linear rule to Adam. The rule was derived and tested for SGD with momentum. Adaptive optimizers already normalise by gradient magnitude, and a square-root rule, lr × sqrt(world), is the more common starting point there. Both are heuristics — measure.

Scaling weight decay too. Decay is per-step shrinkage, unrelated to gradient noise. Leave it alone, and be aware that a larger rate with plain Adam changes the effective decay; AdamW exists partly to break that coupling.

Believing the rule scales forever. Beyond some batch size the returns collapse, and that size depends on the dataset and the model. Goyal et al. held it to a batch of 8192 on ImageNet, and stated plainly that it broke above that.

Try it yourself

Run gradient_noise.py with batches (64, 256, 1024). Predict the ratio between consecutive rows before you look. Then change 0.5 * in the label noise to 0.05 * and see how the whole column shrinks.

What to learn next

Researcher — Mathematics and papers.

The identity, then the heuristic

With per-rank batch $b$, world size $R$, and per-rank loss defined as a mean over its own $b$ samples, DDP's averaged gradient is

$$ \bar{g} = \frac{1}{R} \sum_{r=1}^{R} \frac{1}{b} \sum_{i \in S_r} \nabla \ell(x_i) = \frac{1}{Rb} \sum_{i \in S} \nabla \ell(x_i) $$

where $S_r$ is rank $r$'s slice and $S = \bigcup_r S_r$ is their disjoint union. That is an exact identity, verified by the first script: the distributed gradient is the large-batch gradient.

For an unbiased minibatch estimator with per-example gradient covariance $\Sigma$, the estimator covariance is $\Sigma / (Rb)$, so the per-coordinate standard deviation falls as $(Rb)^{-1/2}$ — the $\sqrt{4} = 2$ pattern in the second script's output.

The linear scaling rule (Krizhevsky 2014; Goyal et al. 2017) sets $\eta \propto Rb$. Its justification is not the noise scaling but a trajectory argument: $k$ SGD steps of size $\eta$ on batches of size $b$ approximate one step of size $k\eta$ on a batch of size $kb$, provided the gradient does not change much across those $k$ steps. That premise fails exactly where the weights move fastest — the first epochs — which is why warmup is part of the rule rather than an optional extra. Goyal et al. reproduced single-worker ImageNet accuracy up to batch 8192 with a 5-epoch gradual warmup, and reported degradation beyond it.

An alternative derivation from the SDE view of SGD, which holds the ratio $\eta / B$ fixed to preserve the noise scale, gives $\eta \propto \sqrt{B}$ for constant stationary distribution rather than constant trajectory. Both rules appear in practice; the square-root form is the usual starting point for adaptive optimizers, and LARS/LAMB replace global scaling with per-layer trust ratios to push batch sizes into the tens of thousands.

McCandlish et al. (2018) make the ceiling quantitative through the gradient noise scale $B_{\text{noise}} = \operatorname{tr}(\Sigma) / |g|^2$: below it, doubling the batch nearly halves the step count; above it, extra batch buys almost nothing. It is measurable during training from the variance between per-rank gradients, which distributed runs already have on hand.

What to hold fixed when comparing

Comparing world sizes honestly means fixing the number of examples seen, not the number of steps — otherwise the larger world silently gets more epochs. Keep the LR schedule expressed in epochs, and note that BatchNorm statistics are computed per rank by default, so the normalisation batch stays at $b$ even as the optimization batch grows to $Rb$. SyncBatchNorm removes that discrepancy at a communication cost, and its absence is a common reason small-per-rank-batch runs behave unlike their single-GPU baseline.

References

  • Goyal et al. (2017), Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour — the linear scaling rule, gradual warmup, and the batch-8192 ceiling.
  • McCandlish et al. (2018), An Empirical Model of Large-Batch Training — the gradient noise scale and where the returns stop.
  • You et al. (2020), Large Batch Optimization for Deep Learning: Training BERT in 76 minutes — LAMB and layerwise trust ratios.
  • Smith et al. (2018), Don't Decay the Learning Rate, Increase the Batch Size — the same trade seen from the other direction.

What to learn next