Data Parallelism
Every worker holds the full model and a different shard of the data; each computes gradients on its shard and the gradients are averaged. It is the simplest split, and it silently multiplies the batch size.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
The model fits on one device and the data does not — how do several copies of the same model train as one, and what changes about the optimisation?
A speech team's acoustic model trains on ten thousand hours of audio. One GPU gets through an epoch in a week. They have been given four GPUs on one node and need the run in two days, and they want the same model they had, only sooner.
Copy the model to four GPUs, give each a quarter of the data, keep the per-device batch of 64, average the gradients after every step, and update all four copies identically. Same code, four times the samples per second.
The four-GPU run finishes an epoch four times faster and its validation loss is worse than the single-GPU run at every epoch. The per-device batch of 64 became an effective batch of 256: the model now takes a quarter as many optimizer steps per epoch, each with a less noisy gradient.
- The four-GPU run finishes an epoch four times faster and its validation loss is worse than the single-GPU run at every epoch. The per-device batch of 64 became an effective batch of 256: the model now takes a quarter as many optimizer steps per epoch, each with a less noisy gradient.
- Raising the learning rate four times to compensate — the linear scaling rule — makes the first epoch diverge, because the tuned warmup was for the smaller rate and the early steps at four times the rate leave the well-behaved region (Batch Size and Learning Rate).
- A worker occasionally reads one example twice when the shard boundaries do not divide the dataset evenly and the loader pads the last batch. Its gradient is very slightly biased toward the padded examples every epoch; nobody notices for a month.
- Batch normalisation statistics are computed per device over 64 examples, not over the effective 256, so the model trained with four workers has subtly different normalisation behaviour from the single-GPU one (Normalisation Layers).
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- The surrounding model maps a window of audio to a distribution over phonemes; the loss is cross-entropy against a forced alignment. This lesson is about producing the same weights with four workers that one worker would have produced.
- "The same weights" means: the same optimisation trajectory in expectation, converging to a model of the same validation quality after the same number of passes over the data.
- One example is a few seconds of audio with its aligned transcript. The dataset is on local NVMe and the loader is fast enough that one GPU is compute-bound.
- Each worker is assigned a disjoint shard by taking every fourth example from a shuffled index; the shuffling seed is shared so the shards are complementary rather than overlapping.
- The single-GPU run used a batch of 64 with a learning rate and warmup tuned for it over several weeks of experiments.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- Every worker holds an identical copy of the parameters. On each step, worker *k* draws a batch from its shard, runs forward and backward, and produces a gradient
g_k. The workers then compute the mean(g_1 + … + g_N) / Nand every worker applies the same update to its identical copy. Because the update is identical, the copies never diverge. - The mean of the per-worker gradients is exactly the gradient of the mean loss over the union of the per-worker batches. So a step on N workers with batch B each is mathematically one step on a single device with batch N·B — not N steps. The number of optimizer updates per epoch drops by N.
- Fewer, less noisy steps change the optimisation. A larger batch estimates the true gradient more precisely, which permits a larger learning rate up to a point, but past that point the gain saturates and the run converges to a different — often flatter, sometimes worse — region in the same number of epochs.
- The averaging is the synchronisation point. Every worker blocks until every other worker's gradient has arrived, so the step runs at the pace of the slowest worker, and the exchange volume is the full parameter count every step regardless of batch size (Gradient Synchronisation).
The averaged gradient is one big step
The identity that makes data parallelism work is the identity that changes the optimisation. The gradient of the mean loss over a batch is the mean of the per-example gradients, so the mean of four workers' gradients on 64 examples each is exactly the gradient on the 256-example union. The four copies apply that same gradient and remain identical.
Four workers therefore take one step of batch 256 in the time one worker takes a step of batch 64. Samples per hour go up four-fold; steps per hour do not move; steps per epoch fall four-fold. Every property of the run that depends on step count — warmup length, schedule, the noise in the gradient — has changed.
1def data_parallel_step(model_copies, shards, lr):2 # each worker: forward + backward on its own shard batch3 grads = [copy.grad(loss(copy, next(shard))) for copy, shard in zip(model_copies, shards)]4 5 # all-reduce: the mean of the per-worker gradients.6 # This equals grad(mean loss over the UNION of the shard batches).7 n = len(grads)8 mean_grad = sum(grads) / n9 10 # every copy applies the same update, so the copies stay identical11 for copy in model_copies:12 copy.params -= lr * mean_grad13 # one step on N*B examples — not N steps on BThe line that matters is the last comment. The per-device batch is an implementation detail; the effective batch N*B is what the optimizer experiences, and the learning rate and warmup were tuned for something else.
The shards and what they must promise
The identity above assumes the union of the shard batches is a fair sample of the data. It is a fair sample when the shards are disjoint, jointly complete, and shuffled with a shared seed so that each worker's batch is as random as a single device's would have been.
The ways shards fail are all quiet. A padded remainder duplicates a few examples per epoch; a sorted file listing gives each worker a biased slice; a worker that dies takes its shard out of the epoch. Each biases the average a little, and a little bias per step over a million steps is a different model.
looks like Four workers each reading dataset[k::4] from a shuffled index, which looks disjoint and complete.
why it leaks It does not leak the label; it double-counts. When the loader pads the last batch of each shard by wrapping to the start, a handful of examples appear twice per epoch on every worker, and their gradient contribution is doubled every epoch for the whole run.
fix Drop the remainder deterministically or distribute it round-robin, and assert per epoch that the count of distinct example ids consumed across workers equals the dataset size.
Re-tuning for the batch you actually have
The learning rate that was tuned for batch 64 is the wrong learning rate for batch 256. A larger batch gives a more precise gradient, which tolerates a larger step; the linear scaling rule multiplies the rate by the batch factor and lengthens the warmup so the early large steps do not diverge. It is a first guess that a short comparison run confirms or corrects.
What cannot be skipped is the comparison itself. A data-parallel run whose validation curve at equal epochs matches the single-device run is the same training, faster. One whose curve is worse is a different, cheaper model, and the decision to accept it belongs to whoever owns the product metric.
The learning-rate schedule in use was tuned for the effective batch — per-device batch times worker count — that the run is actually using.
holds when Batch and worker count are recorded with the run, the schedule was re-tuned when they were last changed, and a short equal-epoch comparison against the single-device baseline was green.
breaks when Someone adds workers to speed a run up, shrinks the per-device batch to fit a larger model, or reuses a schedule from a run with a different worker count — none of which raises an error.
respond Treat the effective batch as a hyperparameter that was changed. Re-tune rate and warmup on a short run before spending the long one.
Effective batch silently multiplies; steps per epoch fall; the run is faster and slightly worse and nobody knows why.
The optimisation problem is chosen on purpose; the schedule matches it; the equal-epoch comparison against the single-device run says whether it is the same model.
The optimizer sees the effective batch, not the per-device one. Deciding the effective batch first makes the worker count an infrastructure choice instead of a hidden hyperparameter.
How to build it
Most important first.
- Decide the effective batch first, then divide by the worker count to get the per-device batch. Keeping the per-device batch fixed and adding workers is a decision to change the optimisation problem, and it should be taken on purpose.
- Re-tune the learning rate and warmup for the effective batch. The linear scaling rule with a longer warmup is a good starting point for many networks; it is a starting point, and the validation curve against the single-device baseline is the test.
- Make the sharding exact: every example index appears on exactly one worker per epoch, drop or distribute the remainder deterministically, and re-shuffle with a per-epoch seed shared by all workers.
- Decide what batch-norm statistics mean across workers — synchronised batch norm, or a per-device batch large enough that it does not matter — and record the decision with the run (Experiment Tracking).
- Verify data coverage and gradient equivalence on a tiny run before the long one: N workers on a fixed batch should produce the same averaged gradient, to floating-point tolerance, as one worker on the concatenated batch.
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- Validation loss against *epochs*, not against wall-clock or steps, compared with the single-device baseline. Equal epochs is the comparison that isolates the effect of the larger batch.
- Samples per second per worker against the single-device rate — the scaling efficiency, which says how much of the four-fold compute the synchronisation is eating.
- A per-epoch coverage count: the number of distinct example ids consumed across all workers, which should equal the dataset size.
- Steps per epoch is not a target. It falls by N by construction; what matters is where the model ends up after the epochs you can afford.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The shards are disjoint and jointly cover the dataset every epoch, with the same shuffle on every worker, so the averaged gradient is an unbiased estimate of the full-data gradient.
- The learning-rate schedule in use was tuned for the current effective batch — the per-device batch times the current worker count — and is re-tuned whenever either changes.
- Every worker starts each step with identical parameters; any divergence between copies (a skipped update on one worker, a mixed-precision mismatch) is detected rather than allowed to accumulate.
- Offline: on a fixed batch, compare the averaged gradient from N workers against the single-device gradient on the concatenated batch. Agreement to floating-point tolerance proves the plumbing; nothing else does.
- Before the long run: a short run with the chosen effective batch and schedule against the single-device baseline on the same subset, compared at equal epochs.
- During: a periodic checksum of the parameters across workers, so a diverged copy is caught at the step it diverges rather than at the end of the run.
What can go wrong
- The effective batch is scaled to a size where the gradient estimate is so precise that the run stops benefiting from the noise — it converges fast to a sharper minimum that generalises worse than the noisier single-device run.
- A re-tuned learning rate that works for four workers is reused when the team moves to sixteen, and the run diverges in the first hundred steps.
- The shards are built from a sorted file listing, so worker 0 gets every speaker whose name starts with A. Each worker's gradient is biased toward its speakers and the average is fine — until a worker dies and its shard vanishes from the epoch.
- The per-worker batch is reduced to fit a larger model into memory, and the compute per step shrinks until the all-reduce dominates. Adding workers now makes the run slower.
- The full parameter vector is exchanged every step. For a model with hundreds of millions of parameters that is hundreds of megabytes per step per worker, which bounds the speed-up on any interconnect.
- The optimisation problem changes with every worker count. A tuned single-device schedule is a starting point, not a result, and the re-tuning is a hyperparameter search that costs its own GPU-hours (The Tuning Budget).
- Every worker holds the whole model, so data parallelism cannot help when the model does not fit; that needs the harder splits.
- "Four workers, four times the steps per hour." Four workers, four times the samples per hour and the *same* number of steps per hour — each step is now a batch four times larger. The step count per epoch falls by four.
- "We kept the per-device batch the same, so nothing about the training changed." Everything about the optimisation changed. The effective batch is what the optimizer sees.
- "Larger batches are always more efficient on GPUs, so scale the batch as far as memory allows." Hardware efficiency and optimisation efficiency are different curves; past some batch size the model needs more epochs to reach the same loss, and the hardware gain is spent on extra passes.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- GENERALThat averaging per-shard gradients equals one step on the union batch follows from the linearity of the gradient of a mean, and holds for any model trained by gradient descent with a per-example loss.
- MODEL-SPECIFICThe sensitivity to the effective batch is model-dependent: convolutional and transformer networks tolerate large batches with a scaled, warmed-up learning rate, while small networks and some recurrent models lose quality quickly as the batch grows, and the safe scaling factor has to be measured per model.
- CONTESTEDThe linear scaling rule — multiply the learning rate by the batch-size factor — is treated by some teams as a law and by others as a heuristic that fails outside the regime it was measured in. The strongest case against it is that it presumes a small-batch, SGD-with-momentum setting, and that adaptive optimizers and very large batches follow a square-root or flatter relationship; the strongest case for it is that it is a good first guess that a short comparison run corrects.
Where the depth lives
This domain teaches the model and hands the rest off by name.