Distributed Training from Scratch — Data Parallel, Model Parallel, and When Communication Becomes the Bottleneck
Why one machine is not enough, counted out in bytes; data parallelism and all-reduce; what ZeRO and FSDP actually shard; tensor and pipeline parallelism. Then the ratio of computation to communication that tells you where scaling stops paying — and gradient accumulation, NCCL settings and how to diagnose a hang.
An analogy: a hundred people revising one dictionary
A hundred people are revising a single dictionary. Splitting up the writing is easy. What is hard is the reconciliation: unless everyone's edits reach everyone else, the dictionary ends up contradicting itself. Hold a full meeting after every entry and the meetings outlast the writing. Add more people and the meetings get longer, until past some headcount another pair of hands makes things worse.
Distributed training has exactly this shape. Splitting the computation is not the difficult part. The difficult part is making every replica agree again afterwards. The protagonist of this article is not the arithmetic unit but the network, and the quantity to estimate is the ratio of compute time to communication time.
Why one machine is not enough
Start by counting bytes. Let be the parameter count and assume the usual mixed-precision setup with Adam. Parameters in bf16 cost bytes and gradients at the same precision another . On top of that sit an fp32 master copy for accurate updates () and Adam's first and second moments ( each).
Put in words: every weight you train drags along a spare high-precision copy of itself and two running notes about its own past, so the bill is charged per parameter rather than per layer. is the memory that stays occupied for the entire run, and is the number of weights.
Sixteen bytes per parameter. At that is roughly 112 GB — the state alone exceeds the memory of a single accelerator before a single number has been computed. And separately from this term, backpropagation needs the intermediate outputs (activations) kept around.
Said in words, activations swell at the same rate whether you make the model deeper, push more samples through at once, or feed it longer text. The terms multiply rather than add, which is what makes this quantity get away from you.
with layers, micro-batch , sequence length and hidden width . It scales with the product of depth, batch and sequence length, so the same model doubles its activation memory when you double the batch. The standard remedy is activation checkpointing: keep only each layer's input, throw the middle away, and recompute during the backward pass. Memory drops sharply, compute rises by roughly one extra forward pass.
So there are two distinct reasons a model does not fit, with two distinct remedies. The of state is reduced by sharding; activations are reduced by recomputation and batch size. Making that distinction first keeps every later decision straight.
Data parallelism and all-reduce
The most natural thing to split is the data. Each of machines holds a full copy of the model, runs forward and backward on its own mini-batch, and then averages only the gradients so the copies stay identical. That is data parallelism, and while the model fits on one device it is always the first thing to try.
The averaging step is an all-reduce: sum everyone's values, then make sure everyone has the sum. Doing it naively — gather to one machine and broadcast back — throttles on that machine's link, so implementations use a ring all-reduce. Arrange the machines in a ring, cut the data into chunks, and make two passes: one accumulating as chunks travel (reduce-scatter) and one circulating the finished chunks (all-gather). The bytes each machine sends are
Put in words: every machine puts about two copies of the whole gradient onto the wire per step, and that is the end of its obligation. counts the bytes one machine sends in one step, not the traffic summed over the cluster.
where is the bytes per gradient element. As grows this only approaches — it is essentially independent of machine count. Thirty-two machines or a thousand, each one ships the same volume. That is why ring all-reduce became the default (the number of hops does grow with , so latency still matters).
The ratio of computation to communication
Where scaling stops paying is decided by comparing one step's compute time against its communication time. As derived in the cost of matrix multiplication, training costs about FLOPs per token. Let be the tokens one machine processes per step and its sustained throughput; compute time is and communication time is from above. The ratio:
The fraction, in words, asks how many times over the arithmetic of a step covers that step's network traffic. Comfortably above one and the transfers disappear behind compute; below one and every machine you add buys you more waiting.
cancelled. Growing the model raises computation and communication in the same proportion, so the balance does not move. Only three things matter: tokens per machine , interconnect bandwidth , and gradient precision .
Comments
Sign in to comment