The previous article measured a single training step: the FLOPs it costs, the bytes it holds, and where its time goes on one device. This one asks the question that follows. Given that workload, what can we split across GPUs, and what communication does each choice force us to pay for?
That question has a finite answer. There are six things in a training step that can be partitioned, no more, and each partition produces exactly one family of techniques. Everything else — the acronyms, the framework flags, the 3-D and 5-D configurations — is a combination of those six with a topology attached.
Distributed vocabulary
Six words first, because every sentence below uses them.
- Rank — a process's global index in the job,
0 … world_size−1. One rank almost always owns exactly one GPU. - World size — how many ranks there are in total. Not how many nodes, and not how many GPUs you reserved: how many processes actually joined.
- Local rank — the rank's index within its node, which is what selects the device (
cuda:local_rank). Confusing local rank with global rank is the classic first-day bug: every process on a node grabs GPU 0 and eight processes fight over one card. - Process group — a named subset of ranks that can perform collectives together. The default group is everybody; the interesting ones are the subsets corresponding to each parallelism axis.
- SPMD — single program, multiple data. Every rank runs the same script; the only thing that differs is its rank number, and therefore which slice of the world it is responsible for. There is no coordinator process, no master program, no task queue.
- Rendezvous — the bootstrap handshake by which ranks discover each other, agree on the world size, and build the default process group. Until it completes, no rank knows it is part of anything.
The SPMD point deserves emphasis, because the discipline of distributed training follows from it. Since every rank runs the same program, correctness depends on every rank reaching the same collective calls in the same order. Nothing enforces this. There is no scheduler checking. If one rank takes a different branch the job does not crash — it stops, silently, until a watchdog kills it.
One program, many ranks
The six partitioning axes
A training step is a dataflow graph. Distributing it does not change the mathematics; it cuts that graph, puts the pieces on different devices, and inserts communication wherever a dependency crosses the cut. So "what parallelism should I use?" is really "where should I cut?", and the list of places you can cut is short:
Communication, in this frame, is not a design choice. It is a consequence. Nobody decided that tensor parallelism should have an AllReduce; the AllReduce exists because cutting a matrix multiply leaves partial sums that no single device can finish. Choose the cut and you have chosen the collective, its size, and its frequency.
Figure 02 runs one step through all six cuts and answers, for each, the questions that determine everything else.
One step, six cuts
Eight questions per strategy
To keep the comparison honest, every strategy below carries the same specification card. The eight entries are not descriptive — each is derivable from the cut, and together they are the cost model:
- What is replicated? Whatever the cut left whole on every device.
- What is sharded? Whatever the cut divided.
- What memory does it save? The sharded part divided by the degree, and nothing else.
- What computation remains? Whether each device still does the full arithmetic or only a fraction of it.
- What communication appears? The collective needed to repair the dependency the cut severed.
- How frequently? Once per step, once per layer, or once per operation. This matters more than volume.
- What topology does it prefer? Frequent and blocking wants a scale-up fabric; infrequent and overlappable tolerates scale-out.
- What scaling limit appears? Every axis runs out, and they run out for different reasons.
The frame for the whole article
Distributed training is not "more GPUs". It is deciding where every tensor lives, when it moves, who waits for it, and whether the useful compute on either side of the cut is large enough to hide the movement.
The collectives, briefly
The communication vocabulary comes from MPI and predates deep learning by decades. The ring algorithms and their cost model were derived in the previous article on cluster communication; this is the refresher needed to read the rest of this one.
- AllReduce — combine elementwise across all ranks, result everywhere. The workhorse of data parallelism.
- AllGather — every rank ends holding the concatenation of all shards. How FSDP reassembles a layer's weights.
- ReduceScatter — combine, then keep only your own chunk. How FSDP produces a sharded gradient.
- All-to-All — a full transpose: rank i's chunk j goes to rank j. What MoE routing needs.
- Broadcast / Reduce / Scatter / Gather — the rooted operations, used at setup, logging and checkpoint time.
- Send / Recv — point-to-point, the only asymmetric primitive. What pipeline stages use.
Two facts carry over and get used constantly below. First, the identity AllReduce = ReduceScatter + AllGather, which is why sharding gradients turns out to be free while sharding parameters is not. Second, the cost model:
The collectives, as shapes
Often said"A collective is a function call, so I can put one inside an if."
ActuallyA collective is a group operation: it completes only when every rank in the group enters a compatible call, with matching shapes and dtypes, in the same order. Guarding one behind rank-dependent control flow is the classic recipe for a hang. This constraint is why distributed code is written in such a rigid, branch-free style, and it is worth internalising before writing any of it.
Data parallelism, derived
Start with the simplest possible cut and let the collective fall out of it.
Give every GPU an identical copy of the model and a different mini-batch. Rank 0 gets batch 0, rank 1 gets batch 1, and so on. Each runs an ordinary forward and backward pass — nothing about the model computation changes, which is what makes this the cleanest starting point.
After the backward pass every rank holds a gradient. But they are different gradients, because they came from different data. If each rank now applied its own update, the replicas would diverge after one step and you would have N independently drifting models rather than one model trained on N times the data.
So: how do we make every replica take the same optimizer step? They must agree on the gradient. And the gradient they should agree on is the one the whole batch would have produced:
Because the averaged gradient is identical on every rank, and every rank started from identical parameters with identical optimizer state, the update is identical too. The replicas stay bit-consistent with no further communication. This is worth stating precisely: DDP is not an approximation of large-batch training. It is large-batch training, computed in pieces. The only sources of difference from a single-device run of the same global batch are floating-point summation order and any operation that couples examples within a batch — which is why BatchNorm needs SyncBatchNorm and LayerNorm does not.
What the global batch is actually made of
Three independent knobs multiply into the number papers report, and conflating them causes real confusion:
no_sync() it skips the intermediate reductions entirely, genuinely lowering communication per token. Under pipeline parallelism a fourth factor appears, because the microbatch count that keeps the bubble small also multiplies the batch.
Underneath that arithmetic sits a correctness question that is easy to get wrong and nearly impossible to see in a loss curve: which examples each rank gets. A distributed sampler must partition the dataset so the union across ranks is the epoch and the intersection is empty. Get the seed handling wrong and two ranks draw the same shard — you are then training on half your data at double weight, and every metric looks normal. The same coordination applies to dropout streams, to sequence packing, and to resuming mid-epoch after a crash.
Where a global batch comes from — and how a sampler breaks it
Data parallelism · specification
- replicated
- parameters, gradients, optimizer state — a full copy on every device
- sharded
- the batch, and nothing else
- memory saved
- none of the model state; only the activations of the samples you no longer hold locally
- computation
- full model, on 1/N of the data per device
- communication
- AllReduce of the gradient buffer, ~2× the gradient bytes per rank
- frequency
- once per optimizer step, spread across the backward pass
- prefers
- anything — it is the most latency-tolerant axis, because it overlaps with an entire backward pass
- limit
- model state must fit on one device; and the global batch eventually exceeds the critical batch size
DDP internals
The interesting part of DDP is not "there is an AllReduce". It is when the reduction starts.
A naïve implementation waits for the entire backward pass, then reduces one giant buffer. That serialises compute and communication: the network idles through the whole backward pass and then becomes the only thing running.
But backpropagation produces gradients in reverse layer order, and the last layer's gradient is final the moment it is computed. It can start moving while earlier layers are still computing. PyTorch DDP exploits this with autograd hooks: each parameter gets a hook that fires when its gradient is ready, parameters are grouped into buckets (25 MB by default, in reverse registration order), and when every parameter in a bucket has reported, that bucket's AllReduce launches asynchronously on a separate stream.
Bucketing solves two problems at once. It overlaps communication with compute, and it fuses many small tensors into few large messages — which, per the α–β model, is the difference between latency-bound and bandwidth-bound.
Gradient buckets firing during the backward pass
A few consequences worth knowing:
- Bucket size is a real tuning knob. Too small and you pay α on every fragment; too large and the first bucket cannot fire until most of the backward pass is done.
- Gradient accumulation with
no_sync()suppresses the hooks for all but the final microbatch, so K microbatches cost one reduction instead of K. It is mathematically identical to the un-accumulated version. - Unused parameters break bucketing. If a parameter's gradient is never produced — a conditional branch, a head that is inactive this step — its bucket never completes.
find_unused_parametersfixes it at the cost of a graph traversal every iteration; a static graph is faster when the model genuinely does not change shape. - Stragglers become everybody's problem. An AllReduce completes when its last participant arrives, so a slow data loader or a throttling GPU on one rank sets the pace for the entire group. We return to this in §Scaling efficiency.
Often said"DDP synchronises gradients at the end of the backward pass."
ActuallyIt synchronises throughout it. The last bucket to fire is the one holding the first layer's parameters, and only that final reduction is genuinely exposed. If you profile DDP and see one large AllReduce at the end, your bucketing has broken — usually unused parameters, a conditional branch, or checkpointing interacting badly with hook ordering.
Overlap is not free
The timeline above draws communication on its own lane, as though issuing a collective on a separate CUDA stream made it cost nothing. It does not. NCCL collectives are kernels: they occupy streaming multiprocessors, contend for memory bandwidth, and compete with your matmuls for copy engines. A reduction running "in parallel" with a GEMM makes that GEMM slower.
So the honest model has three resources, not two. A collective is genuinely hidden only when the network can absorb the bytes in the time available and the SMs it borrows were ones the compute would not have used. On a well-fed GEMM that is rarely exactly true, which is why measured overlap efficiency usually lands between 60% and 90% rather than the 100% every timeline diagram implies.
The three resources a "free" overlap actually consumes
Often said"Communication is fully overlapped, so it is free."
ActuallyOverlap moves the cost from the timeline into the throughput. The step gets no longer, but the compute inside it runs slower because the collective is holding SMs and memory bandwidth. Compare MFU with and without the collective rather than comparing step times, and the tax appears immediately. Budget 5–15% on a busy step.
Why DDP stops being enough
Two separate failures, and it matters which one you have, because they call for different medicine.
Failure one: memory
Everything except the batch is replicated. Parameters, gradients, optimizer moments and master weights all sit on every device in full. Adding GPUs multiplies your aggregate memory but leaves per-device memory exactly where it started — so a model that does not fit on one card does not fit on a thousand of them either. This is not a bandwidth problem or a tuning problem; it is a hard wall, and no amount of DDP crosses it.
Adding GPUs does not add memory
Failure two: scaling
The second failure is subtler, and it is where strong and weak scaling part company.
Weak scaling holds the per-device batch fixed and adds devices, growing the global batch. Compute per device is constant, so efficiency erodes only through the exposed part of communication. This scales well, and it is what nearly every published scaling chart shows.
Strong scaling holds the global batch fixed and divides it further. Now per-device compute shrinks as 1/N while the gradient AllReduce does not shrink at all — its size is set by the parameter count, not the batch. The ratio of communication to compute therefore grows linearly with N, and efficiency collapses far earlier than weak scaling would suggest.
Both failures point the same way: stop replicating things.
Deriving ZeRO
Here is the observation that starts it. In DDP, every rank stores an identical copy of the optimizer state, the gradients and the parameters. Identical. That is N−1 copies of pure redundancy, held for no reason except convenience — and a rank only needs the full parameter tensor at the instant it is computing with it.
So ask: why replicate things no device needs simultaneously? Remove the redundancy one term at a time and you get the three ZeRO stages. Using the mixed-precision Adam ledger from the previous article — 2 bytes of bf16 parameters, 2 of bf16 gradients, and 12 of fp32 master weights plus moments, so 16 bytes per parameter Ψ:
The communication column is the point. Stages 1 and 2 are volume-neutral: ReduceScatter followed by AllGather is precisely what an AllReduce does internally, so they cost nothing extra. Sharding optimizer state and gradients is free memory, and there is essentially never a reason to run without it.
Stage 3 is different. Because the weights were freed after the forward pass, they must be gathered a second time during backward — an extra Ψ that neither of the others pays. That is the 1.5× factor, and it is the entire price of the memory below it.
The ZeRO ladder, one shard at a time
Often said"ZeRO-3 has the same communication as ZeRO-2, because ReduceScatter + AllGather = AllReduce."
ActuallyThat identity holds for the gradient path, which is why stages 1 and 2 are volume-neutral. Stage 3 adds a second parameter AllGather in the backward pass that neither of the others needs. You can trade it back with reshard_after_forward=False, spending memory to buy bandwidth.
Sharded data parallelism (ZeRO) · specification
- replicated
- the computation — every device still executes every layer
- sharded
- optimizer state (Z1), plus gradients (Z2), plus parameters (Z3)
- memory saved
- 12Ψ→12Ψ/N, then 14Ψ/N, then the whole 16Ψ/N
- computation
- unchanged — full model, 1/N of the data
- communication
- ReduceScatter + AllGather; stage 3 adds a second parameter AllGather
- frequency
- once per step for Z1/Z2; once per sharded unit for Z3
- prefers
- Z1/Z2 anywhere; Z3 wants the shard group inside a fast fabric
- limit
- memory saving is 1/N with sharp diminishing returns, while per-layer collectives get slower as the group widens
FSDP
ZeRO-3 as a memory equation is one thing; ZeRO-3 as an execution schedule is another. FSDP is that schedule. It groups parameters into units — typically one transformer block each — and gives each unit a rhythm:
Several policy knobs follow directly from this rhythm, and each has a characteristic failure:
- Wrapping granularity. The unit size is the communication granularity. Per-parameter wrapping produces messages too small to saturate a link and drowns you in α; whole-model wrapping defeats sharding entirely. One transformer block per unit is the default for good reason.
- Prefetch. The AllGather for layer
i+1depends only on shards that are already resident — never on layeri's output. So it can be issued while layericomputes. With correct prefetching the 1.5× volume costs nothing in wall-clock; without it every layer stalls on its own weights and FSDP looks catastrophically slow. A backward prefetch policy ofBACKWARD_POSTinstead ofBACKWARD_PREissues the gather too late to hide. - Reshard after forward. Keeping the gathered parameters resident through the backward pass removes the second AllGather, converting the 1.5× back to 1×, at the cost of holding a full layer's weights. Usually correct for the last stage of a run, where memory pressure has eased.
- Memory spikes. Peak is not average. The moment a unit is gathered, that unit exists in full on top of everything else. A single oversized unit — a tied embedding, a large LM head — can OOM a configuration whose average footprint is comfortable.
- CPU offload moves optimizer state and gradients to host memory, trading PCIe bandwidth for HBM. It makes otherwise impossible runs possible and is almost always slower; treat it as a capacity tool, not a performance one.
FSDP with and without prefetch
Hybrid sharding: the shard group is a choice
FSDP's cost depends entirely on the bandwidth of the group you shard across. Inside a node, 8 GPUs share NVLink/NVSwitch at hundreds of GB/s. Between nodes, even a good fabric gives an order of magnitude less. Sharding across all 512 GPUs of a cluster puts the per-layer AllGather on the slow fabric — exactly where it cannot be hidden.
HSDP arranges the world as a 2-D mesh: shard within a node, replicate across nodes. The frequent, latency-critical per-layer collectives ride the fast fabric, while cross-node traffic becomes a single gradient AllReduce per step, which overlaps with the whole backward pass. You keep 1/8 of the memory saving instead of 1/512 — but if the model fits at 8-way sharding, the other 504 GPUs are better spent on throughput.
HSDP: putting the chatty collective on the fast wire
Often said"FSDP shards everything, so memory scales as 1/N."
ActuallyFSDP shards persistent state. Activations are not sharded — they belong to the local microbatch and stay where they were produced. On long-context or large-microbatch runs activation memory dominates and FSDP does nothing for it. That gap is precisely what checkpointing and context parallelism exist to fill, and why "I enabled FSDP and still OOM'd" is such a common report.
Often said"Shard as widely as possible; more sharding is more memory saved."
ActuallyGoing from 8-way to 64-way sharding takes you from 12.5% of the ledger to 1.6% — eleven percentage points — while moving every per-layer collective onto a wire that may be an order of magnitude slower. Shard to fit, then replicate.
Tensor parallelism
Sharded data parallelism divides storage but not arithmetic: every rank still runs the whole layer, it just borrows the weights first. Tensor parallelism asks a different question — what if a single matrix multiply is too large for one device, or simply too slow? — and answers it by splitting the multiplication itself.
The construction rests on two ways of cutting a linear layer, and on the fact that they chain so cleanly that a whole transformer block needs only one collective.
Column parallel: split the output features
Row parallel: split the input features
The nonlinearity between them is the subtle part, and it is why the ordering is not arbitrary. GELU is elementwise, so GELU([Y₁|Y₂]) = [GELU(Y₁)|GELU(Y₂)] — it commutes with a column split. Had you used row-parallel first, you would need an AllReduce before the nonlinearity, doubling the collectives. Megatron's layout puts the nonlinearity exactly where the tensor is already correctly sharded.
Attention shards even more naturally, because heads are independent. Assign heads to ranks: the QKV projections are column-parallel (producing per-rank heads), the attention itself is local to those heads, and the output projection is row-parallel. One AllReduce, same as the MLP.
Column-parallel then row-parallel, block by block
Two properties follow immediately. First, TP shards both memory and arithmetic — the only strategy so far that shrinks the matmul itself. Second, its communication sits inside the layer, on the critical path, with nothing to overlap against: the next operation literally cannot begin until the sum exists. That combination is why tensor parallelism lives and dies by its fabric.
The tensor-parallel tax, on the fast fabric and off it
Often said"Set TP = 8 to use the whole node, always."
ActuallyTP shrinks each rank's matmuls. At TP=8 a hidden size of 8192 becomes per-rank GEMMs of width 1024, which start to under-fill the tensor cores and lose efficiency to launch overhead and reduced arithmetic intensity. TP=8 is right for very large hidden sizes; for 7B-class models TP=2 or 4 often beats it, and TP=1 with FSDP frequently beats both. Measure the GEMM, not the theory.
Often said"Tensor parallelism is just splitting the model somehow."
ActuallyThe algebra of the split is chosen so that the resulting matmuls stay large and efficient and the collectives fall on natural block boundaries. A bad cut produces tiny GEMMs, extra materialisation, or a collective after every sub-operation. The column-then-row ordering is not one option among many; it is the one that needs a single AllReduce per block.
Tensor parallelism · specification
- replicated
- the block's input activations, and everything between the parallel regions — layernorm, dropout, the residual stream
- sharded
- each weight matrix along one axis; attention heads across ranks
- memory saved
- weights and inner activations divided by T; the regions between are untouched
- computation
- genuinely divided — each rank does 1/T of the arithmetic
- communication
- 4 AllReduces per layer per step, each of size b·s·h
- frequency
- twice inside every layer, blocking, on the critical path
- prefers
- a scale-up fabric — NVLink/NVSwitch. Effectively node-local
- limit
- per-rank GEMM width, head count, and the node boundary. Rarely above 8
Sequence parallelism
Look again at the TP specification card: everything between the parallel regions is replicated. Layernorm inputs, dropout masks and the residual stream sit on every TP rank in full. For long sequences those regions are not small — and worse, they scale with b·s·h, exactly like the activations TP did shard.
Megatron-style sequence parallelism closes that gap. In the regions where TP has nothing to shard, split along the sequence dimension instead: rank i holds tokens [i·s/T, (i+1)·s/T) of the layernorm input, computes its share, and hands off. The two layouts alternate through the block.
The saving is real: the replicated layernorm, dropout and residual activations drop by a factor of T. Since these are the terms that survive TP untouched, SP is what makes tensor parallelism's memory story match its arithmetic story.
Alternating layouts: TP regions and SP regions
Often said"Sequence parallelism and context parallelism are two names for the same thing."
ActuallyThey are routinely confused and solve different problems. Megatron's sequence parallelism is an add-on to tensor parallelism, riding on TP's degree, sharding the regions TP leaves replicated, and changing no collective count — only their shape. Context parallelism is a parallelism axis in its own right, with its own degree and its own process group, and it requires a genuine distributed attention algorithm because local queries must reach remote keys. A run can have both, and large long-context runs usually do.
Sequence parallelism · specification
- replicated
- nothing new — it removes replication rather than adding it
- sharded
- layernorm, dropout and residual activations, along the sequence axis
- memory saved
- the non-TP activation regions, divided by T
- computation
- the elementwise regions divided by T; matmuls unchanged
- communication
- none added — each AllReduce becomes ReduceScatter + AllGather
- frequency
- the same boundaries tensor parallelism already had
- prefers
- whatever the TP group prefers; it shares that group
- limit
- bounded by the TP degree — it has no degree of its own
Context parallelism
Every axis so far shards parameters or batches. Long-context training breaks both, because the offending quantity is neither: it is the activation memory of a single sequence, and it grows with s.
Consider one sequence of 128k tokens. Per-layer activations scale as b·s·h, which at s=131072 and h=4096 is around a gigabyte per layer per sample in bf16 before attention. You cannot reduce the batch below one. FSDP does not shard activations. TP replicates the residual stream and is capped at the node. Pipeline parallelism divides depth, not sequence. The only remaining axis is the sequence itself.
Context parallelism shards the token dimension across C devices. Everything token-local — layernorm, MLP, elementwise ops — becomes trivially parallel with no communication. Attention is the entire difficulty, because token i attends to every token before it, including those on other devices.
Ring attention
Keep queries local; rotate keys and values around a ring. At step k, each rank holds the KV block belonging to rank (i−k) mod C, computes the partial attention of its local queries against that block, and accumulates. After C steps every query has seen every key. Communication is one KV block per step, prefetchable during the previous block's compute — so with enough tokens per rank it hides completely.
Causal masking makes the work uneven: rank 0's queries need only one KV block, rank C−1's need all C. Implementations either load-balance by giving each rank two non-contiguous chunks — one from each end of the sequence — or accept a 2× imbalance.
Ulysses
A different decomposition. Before attention, an All-to-All converts a sequence-sharded, head-complete layout into a head-sharded, sequence-complete one. Each rank then owns a few complete heads over the whole sequence and runs entirely standard attention with no modification. A second All-to-All converts back. Four all-to-alls per layer, each moving b·s·h/C per rank.
The trade is clean. Ulysses moves less data and reuses the unmodified attention kernel, but C cannot exceed the number of key/value heads, and All-to-All is the pattern most punished by an oversubscribed fabric. Ring attention has no head limit and uses only nearest-neighbour point-to-point, but its communication scales with sequence length and its causal balance needs care. In practice: Ulysses inside a node, ring across nodes, often both at once.
Ring attention and Ulysses, side by side
Often said"Attention is O(s²), so long context is hopeless."
ActuallyAttention compute is O(s²) and always was. What FlashAttention removed was the O(s²) memory, by never materialising the score matrix. Context parallelism then divides the remaining quadratic compute across C devices, giving O(s²/C) per rank. Long context is expensive, not impossible — and the expense is now honest FLOPs rather than an allocator failure.
Context parallelism · specification
- replicated
- all parameters and optimizer state
- sharded
- sequence positions, and every activation that follows them
- memory saved
- activations divided by C — the only axis that touches a single sequence
- computation
- divided by C, including the quadratic attention term
- communication
- ring: KV blocks, nearest-neighbour. Ulysses: 4 All-to-Alls per layer
- frequency
- inside attention only; the rest of the block is silent
- prefers
- Ulysses wants a fast fabric; ring tolerates scale-out because it is P2P and prefetchable
- limit
- Ulysses is capped by KV head count; ring is capped by bandwidth per rotation step
Pipeline parallelism
Instead of splitting individual operations, why not assign whole groups of layers to different GPUs? Stage 0 holds layers 1–20, stage 1 holds 21–40, and activations flow forward along the chain while gradients flow back.
That is the enormous advantage: pipeline communication is tiny and infrequent. One tensor of shape [b_micro, s, h] crosses each boundary per microbatch, using point-to-point send/recv, versus tensor parallelism's four full AllReduces per layer. Pipeline stages are the only form of model parallelism genuinely comfortable across a slow network.
The equally enormous disadvantage is the bubble. Run one batch through a 4-stage pipeline and stage 3 idles until stages 0, 1 and 2 have finished; then stages 0–2 idle while the backward pass drains. The fix is microbatching: split the batch into m microbatches and keep several in flight.
1F1B: the same bubble, a fraction of the memory
GPipe's schedule holds all m microbatches' activations alive on stage 0 before any backward begins — m× the activation memory, which usually breaks before the bubble does.
1F1B reorders the same operations: after a warm-up of P−i−1 forwards, stage i alternates strictly between one forward and one backward. The bubble is identical — the same (P−1) idle slots at each end — but live activations per stage drop from m to at most P. Same throughput, memory independent of microbatch count. There is almost no reason to run GPipe scheduling today.
Interleaving goes further: give each device v non-contiguous chunks of layers instead of one contiguous block. The bubble becomes (P−1)/(v·m + P−1), paid for with v× more point-to-point messages and more live activation stashes. Zero-bubble schedules take the last step, splitting the backward into its input-gradient and weight-gradient halves — the weight gradient has no downstream dependency, so it can be deferred into the bubble.
Four pipeline schedules on one clock
Often said"1F1B is faster than GPipe."
ActuallyThey have exactly the same bubble ratio and the same throughput at the same m. 1F1B wins because its activation memory is bounded by P instead of m, which lets you raise m — and raising m is what shrinks the bubble. It is an indirect win, through the memory budget.
Often said"Split the layers evenly across stages."
ActuallyBalance by time, not layer count. The first stage carries the token embedding and the last carries the output projection and loss — for a large vocabulary the LM head plus its softmax can cost more than several transformer blocks, and it materialises a [b·s, V] logit tensor that dwarfs everything else. Give the first and last stages fewer transformer layers, and compute the loss in chunks so the full logit tensor never exists. A pipeline runs at the speed of its slowest stage, and one heavy stage makes every other stage's bubble worse.
Often said"More microbatches is always better, because the bubble shrinks."
ActuallySmaller microbatches mean smaller GEMMs, and below a certain size the matmuls stop saturating the GPU. You trade a bubble you can see in the schedule for an efficiency loss you can only see in MFU. Pipeline tuning is a three-way trade between bubble, kernel efficiency and memory.
Pipeline parallelism · specification
- replicated
- nothing — each layer lives in exactly one place
- sharded
- the layer stack, by depth
- memory saved
- all model state divided by P; activations bounded by the schedule
- computation
- divided by P, but with a bubble of (P−1)/(m+P−1)
- communication
- point-to-point send/recv of one activation tensor per boundary
- frequency
- once per microbatch per stage boundary — by far the least chatty axis
- prefers
- anything, including racks apart. This is the axis to spend slow links on
- limit
- the bubble, stage-balance skew, and the microbatch count needed to hide both
Expert parallelism
Every axis so far assumes every token passes through every parameter. Mixture-of-experts breaks that assumption: replace the MLP in each block with E parallel MLPs and a small router that sends each token to the top-k of them. Parameters multiply by roughly E; FLOPs per token multiply by k, typically 1 or 2.
Which raises the placement question. Why would every device hold all E experts when each token uses only k of them? Put different experts on different ranks. But tokens arrive wherever the data-parallel split left them, so a token on device 3 whose expert lives on device 11 must physically travel there, be processed, and come back. Across the batch, every device sends some tokens to every other device — an All-to-All. Two of them, dispatch and combine, per MoE layer, in each direction.
Routing is learned, and left alone it collapses: a few experts get slightly better early, receive more tokens, improve faster, and absorb everything. MoE training therefore carries an auxiliary load-balancing loss penalising the correlation between routing probability and token fraction. Too little and experts collapse; too much and the router stops specialising.
Token routing, capacity, and the two all-to-alls
Three refinements matter in practice. Top-k changes traffic directly: top-2 doubles both expert compute and token movement relative to top-1. Dropless routing is increasingly preferred — grouped or ragged matmul kernels tolerate uneven expert loads, removing the information loss but handing the imbalance straight to the critical path, since the layer still finishes when the busiest expert does. And experts can have their own tensor parallelism: when a single expert is large, EP chooses which rank group receives a token while expert-TP splits that expert's MLP across the ranks of the group.
Often said"All-to-All moves the same bytes as AllReduce, so it costs about the same."
ActuallyVolume is similar; the pattern is not. Ring AllReduce uses only nearest-neighbour links and can be scheduled so no link contends. All-to-All is all-pairs: on a fat-tree with any oversubscription it drives traffic through the upper tiers from every node simultaneously, and achieved bandwidth can fall well below the nominal per-link number. This is the pattern that most rewards keeping expert groups inside a node.
Expert parallelism · specification
- replicated
- attention, layernorms, and the router itself
- sharded
- the expert MLPs, by expert index
- memory saved
- expert weights divided by E — the dominant term in an MoE model
- computation
- sparse: k of E experts run per token, so FLOPs stay near a dense model of the base size
- communication
- All-to-All dispatch and combine, carrying k·b·s·h per layer per direction
- frequency
- twice per MoE layer, forward and backward
- prefers
- high bisection bandwidth; expert groups ideally inside one node
- limit
- routing imbalance — the layer runs at the speed of the busiest expert, whatever the average load
Combining parallelisms
Real runs use several axes at once, and the mental model that makes this tractable is a device mesh: view your N GPUs not as a list but as a multi-dimensional grid, where each axis is one parallelism strategy and a GPU's coordinate along an axis says which shard of that strategy it holds.
The device mesh, explorable
The dimensions also interact multiplicatively on the batch. Global batch = b_micro × m × DP, where m is the pipeline microbatch count. Raising PP forces you to raise m to keep the bubble small, which raises the global batch, which may push you past the critical batch size. A systems decision changing an optimisation outcome — the part that never appears in throughput charts.
Topology-aware placement
Ranks are not interchangeable. Eight GPUs in a node share a switched, high-bandwidth, low-latency domain; two GPUs in different racks share a fabric that may be an order of magnitude slower and oversubscribed besides. Since each parallelism axis has a characteristic communication frequency and blocking behaviour, matching axes to fabric tiers is most of the performance work.
The general intuition, from the fastest links outward:
- Very frequent, blocking, activation-sized → keep close. Tensor parallelism issues four AllReduces per layer directly on the critical path. It belongs inside the scale-up domain, and effectively never crosses a node boundary.
- Frequent but prefetchable → next. Context parallelism's KV exchange and FSDP's per-layer AllGather can be issued ahead of their consumer. They tolerate one tier down, provided the prefetch distance covers the latency.
- All-pairs → wants bisection bandwidth. Expert parallelism's All-to-All is not latency-sensitive so much as topology-sensitive; it punishes oversubscription specifically. Keep expert groups within a node where possible.
- Infrequent and overlappable → scale-out is fine. Data-parallel gradient reduction happens once per step and overlaps with an entire backward pass. This is the axis to spend inter-node bandwidth on.
- Neighbour P2P, tiny messages → spend the worst links here. Pipeline parallelism sends one activation tensor per microbatch boundary. It will happily cross racks.
These are design tendencies, not laws. A cluster with a genuinely flat, non-oversubscribed fabric relaxes most of them; a cluster with one NIC per node tightens all of them. And two axes that are each efficient in isolation can interfere when they share a NIC — a tensor-parallel AllReduce and a data-parallel reduction issued simultaneously do not each get full bandwidth. Multi-dimensional parallelism is a traffic-engineering problem as much as an algorithmic one.
Often said"3-D parallelism means DP × TP × PP, and the order does not matter."
ActuallyThe order determines which ranks share a node, which determines which collective rides the fast fabric. The same 512 GPUs with the same degrees can differ by well over 2× in step time depending purely on rank assignment. This is why frameworks expose mesh ordering explicitly, and why device-visibility shuffles and bad scheduler rank-to-node mappings are a recurring source of mysterious slowdowns.
Often said"Use every parallelism dimension — more axes, better scaling."
ActuallyEach dimension adds a failure mode, a divisibility constraint, and usually some replicated state. The best configurations are the simplest ones that fit. A 7B model on 64 GPUs should almost certainly be plain FSDP or HSDP with nothing else enabled; adding TP and PP will make it slower, harder to debug, and no more capable.
Overlap, strategy by strategy
Communication volume is the number everyone quotes, and it is the less important one. What determines throughput is how much of that communication can be hidden behind compute — which depends entirely on whether the consumer of the collective's result is the very next operation.
Ranked from most hideable to least:
- DDP gradient reduction. Perfectly overlappable in principle: buckets fire as the backward pass proceeds, and only the final bucket is exposed. The consumer is the optimizer step, which is an entire backward pass away.
- FSDP parameter AllGather. Overlappable with one layer of lookahead, because the gather depends only on resident shards. Requires explicit prefetching to realise; without it, fully exposed.
- Pipeline P2P. Naturally concurrent — while stage i sends microbatch j forward, it is already computing microbatch j+1. The pipeline's inefficiency is the bubble, not the transfer.
- Context-parallel KV exchange. Overlappable by design: the block for rotation k+1 is sent while rotation k computes. Hides well when tokens per rank are large.
- MoE All-to-All. Partially overlappable with careful kernel design, but dispatch must complete before expert compute starts and combine before the block continues. Both sit near the critical path.
- Tensor-parallel AllReduce. Not overlappable at all. The next operation consumes the reduced tensor directly. This is why TP is the axis that most needs a fast fabric, and why its cost shows up in step time rather than in MFU.
This ordering is the practical reason the mesh is laid out the way it is. It is also a useful diagnostic: if a run is slower than the volume arithmetic predicts, the question is not "how many bytes" but "which of these is exposed".
Where each strategy's communication sits
Choosing the strategy
Start from the bottleneck, not from the framework's feature list. Ask what fails on one GPU — model-state memory, activation memory, a single oversized layer, total step time, or long-context attention — and then choose the least communication-heavy axis that fixes it.
The decision tree
Written as a ladder, with the step everyone skips at the top:
- Make the single-device step efficient first. bf16 with fused optimizer kernels, FlashAttention or SDPA, activation checkpointing where it pays, and a dataloader that is not the bottleneck. Distributing an inefficient step just buys many inefficient steps.
- If the model fits, use data parallelism — with ZeRO-1, which is free.
- If model states do not fit, shard them: FSDP/ZeRO-3, at the smallest degree that fits, replicating beyond it.
- If one layer is itself too large, or you need more compute per layer, add tensor parallelism — inside the fast-link domain only.
- If depth still exceeds memory, or you need model partitioning across nodes, add pipeline parallelism.
- If long context dominates activation memory, add context parallelism — after activation checkpointing, which costs predictable compute rather than unpredictable bandwidth.
- If the architecture is MoE, add expert parallelism and design the topology around All-to-All.
| Strategy | What is split | Memory effect | Communication | Sync granularity | Reach for it when |
|---|---|---|---|---|---|
| DDP | the batch | none — 16Ψ replicated | AllReduce, ~2× grad bytes | backward buckets | the model already fits |
| ZeRO-1/2 | optimizer, gradients | 16Ψ → 4Ψ, then 2Ψ | same as DDP | backward buckets | always — it is free |
| FSDP / ZeRO-3 | everything persistent | 16Ψ/N | 1.5× DDP, per layer | per sharded unit | model states do not fit |
| Tensor | each matmul | weights and inner activations / T | 4 AllReduces per layer, activation-sized | sub-layer, blocking | individual layers are huge; stays in-node |
| Pipeline | the layer stack | weights / P | P2P, one activation per boundary | pipeline clock | deep models; slow network |
| Context | the sequence | activations / C | KV ring or 4 all-to-alls | attention layer | activation OOM from long sequences |
| Expert | the experts | expert weights / E | 2 All-to-Alls per layer | MoE routing | mixture-of-experts models |
The same seven strategies, as one moving picture
Figure 21 runs that ladder quantitatively. Describe your model and your cluster and it produces a configuration with the reasoning shown, plus the memory and communication estimates behind it.
The planner lab
Distributed checkpointing
A checkpoint is not a file. It is a distributed data structure, and its shape is determined by the parallelism you chose.
Under FSDP each rank owns a slice of every optimizer tensor. Under TP the slices are cut along a different axis. Under PP whole layers live in different processes. Saving that state means choosing between two options with very different consequences:
- Gathered (full) checkpoints. All ranks contribute to full tensors written by one writer. Trivially loadable at any world size, but slow, memory-hungry at save time, and bottlenecked on a single process's storage bandwidth.
- Sharded checkpoints. Every rank writes its own slice in parallel, so write bandwidth scales with the cluster. Fast — but the files now encode your current mesh.
The consequence is that resuming on a different number of nodes is a resharding operation, not a load. If you wrote sharded checkpoints without enough metadata to redistribute them, you have quietly made your topology permanent — a fact discovered at the worst possible moment, after a node failure, when the replacement capacity has a different shape and the run will not restart. Store the mesh description alongside the tensors and test a reshard before you need one.
Checkpoint geometry and resharding
Checkpointing is also a throughput decision, not just a safety feature. If a node has a mean time between failures of one year, 512 nodes have an expected failure every 17 hours; a month-long run will be interrupted dozens of times. Checkpoint too rarely and you lose hours per failure; too often and you spend a meaningful fraction of the run writing terabytes. The optimum interval is roughly √(2 · T_checkpoint · MTBF) — the classic Young–Daly result, worth actually computing.
Often said"We checkpoint every 1000 steps, which feels safe."
Actually"Feels safe" is not a unit. Compute the interval from your measured checkpoint duration and your observed failure rate; clusters differ by an order of magnitude in both. Also measure the restart cost — process launch, rendezvous, NCCL bootstrap and dataloader warm-up can exceed the write itself.
Scaling efficiency
One definition, used consistently, prevents a lot of self-deception:
Strong scaling fixes the total work and adds devices. Per-device compute falls as 1/N while communication does not, so η degrades quickly. Weak scaling fixes the per-device work, so compute stays constant and only exposed communication erodes η. Large-model training is almost always weak scaling, and almost every near-linear scaling chart you have seen is a weak-scaling chart.
There is a second ceiling with nothing to do with hardware. Weak scaling grows the global batch, and past a model-dependent critical batch size the extra examples stop reducing gradient noise meaningfully. Beyond it you spend twice the FLOPs to reach the same loss. Cluster efficiency can read 95% while training efficiency quietly halves.
Strong scaling, weak scaling, and the critical batch size
Often said"We got 94% scaling efficiency at 512 GPUs, so the system is nearly perfect."
ActuallyScaling efficiency is measured against your own single-node number. If that baseline achieves 30% MFU, then 94% of it is 28% MFU and most of the cluster is idle in a way the scaling chart cannot show. Report MFU alongside η or the number means very little.
Failure modes
The mathematics can be perfectly correct while the system schedule is slow, unstable or deadlocked. These are the failures that actually consume time.
Hangs, not crashes
A crash gives you a stack trace. A hang gives you a job burning cluster hours doing nothing until a watchdog kills it. All the hang modes are variations on one rule: every rank in a process group must call compatible collectives in the same order. A ring reduction cannot complete if a participant never arrives, and NCCL cannot distinguish "still computing" from "never coming".
- Divergent control flow. A data-dependent conditional — skipping a loss term for a batch with no valid targets, an early
continueon a malformed sample, evaluating on rank 0 only — makes one rank call a different collective, or none. - Mismatched collectives. Right operation, wrong buffer: variable sequence lengths and per-rank padding are the usual culprits, and the symptom is often corruption rather than a clean error.
- Wrong process group. With five axes there are five groups. A correct collective on the wrong group sometimes even completes — producing a silently wrong answer rather than a hang.
- A rank that is simply gone. Host evicted, process killed by the OOM reaper, container restarted. Nothing will ever arrive from it.
How a distributed run hangs
The defences are unglamorous and worth building early: compute branch conditions then AllReduce them so every rank takes the same path; assert shapes and dtypes before collectives during bring-up; set a NCCL timeout short enough that a hang fails loudly; name every process group and pass it explicitly; and log each rank's last collective, because comparing four lines answers in seconds what otherwise costs hours.
A single slow rank
A synchronous collective completes when its last participant arrives. If one GPU is 5% slow — thermal throttling, a noisy neighbour on the NIC, an ECC error being corrected — the entire cluster is 5% slow, not 5%/N slow.
The straggler tax
The rest of the list
- OOM during AllGather. FSDP's peak is not its average: the instant a unit is gathered, it exists in full on top of everything else. A configuration with comfortable average memory can die on one oversized unit.
- Topology mistakes. A TP group split across nodes, or a scheduler rank-to-node mapping that interleaves where you assumed contiguity. Silent, and easily a 2× slowdown.
- Pipeline imbalance. One heavy stage sets the clock and widens every other stage's bubble. Embeddings and the LM head are the usual offenders.
- MoE imbalance. A hot expert makes its rank the straggler for the whole layer, every layer, every step.
- Network congestion. Two well-tuned groups sharing a NIC, or All-to-All traffic colliding with gradient reduction. Each is fine alone.
- Small-message death. Fine-grained sharding that reduces memory while converting large bandwidth-efficient transfers into a storm of latency-bound collectives.
Often said"The loss curve looks the same, so the distributed setup is correct."
ActuallyMany bugs are invisible in the loss for hundreds of steps: a sampler giving ranks duplicate samples, a gradient summed instead of averaged (silently scaling your learning rate by N), an unsharded parameter drifting out of sync, a dropped MoE token that should not have been dropped. Test the boring things — assert parameter equality across DP ranks after a step, check the gradient norm against a single-device run on the same data, and verify the number of unique samples per epoch.
Changing the problem
Everything above optimises throughput for a workload that is known in advance. The shapes are fixed, the step is the same every time, the batch is chosen by you, and the only adversary is the network. That is what makes training tractable to reason about: it is a scheduling problem with static inputs.
Serving is not that. There is no backward pass, so two thirds of the FLOPs vanish and most of the memory ledger with them — which sounds easier, and is not. What replaces them is a workload that arrives dynamically, one request at a time, with no advance knowledge of how long any of them will run; and a computation that is autoregressive, so its state grows one token at a time and must be kept alive between steps for every request in flight simultaneously.
The unit of work stops being a step and becomes a token. The memory that matters stops being parameters and becomes KV cache. The question stops being "how do I split this" and becomes "how do I batch things that did not arrive together". That is the next article: The Mechanics of LLM Inference.
References
Grouped by topic. Every entry links to its primary source. Links were written without network access and should be verified before publication.
Data parallelism and sharding
- PyTorch Distributed: Experiences on Accelerating Data Parallel Training · Li et al., 2020 — bucketing, hooks and overlap
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models · Rajbhandari et al., 2019
- PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel · Zhao et al., 2023
- ZeRO-Offload: Democratizing Billion-Scale Model Training · Ren et al., 2021
Tensor, sequence and pipeline parallelism
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism · Shoeybi et al., 2019
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM · Narayanan et al., 2021 — interleaved 1F1B and the 3-D mesh
- Reducing Activation Recomputation in Large Transformer Models · Korthikanti et al., 2022 — sequence parallelism and selective checkpointing
- GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism · Huang et al., 2018
- PipeDream: Fast and Efficient Pipeline Parallel DNN Training · Harlap et al., 2018 — 1F1B scheduling
- Zero Bubble Pipeline Parallelism · Qi et al., 2023
Long context
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness · Dao et al., 2022
- Ring Attention with Blockwise Transformers for Near-Infinite Context · Liu, Zaharia & Abbeel, 2023
- DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models · Jacobs et al., 2023
- Online normalizer calculation for softmax · Milakov & Gimelshein, 2018 — the rescaling identity ring attention depends on
Mixture of experts
- GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding · Lepikhin et al., 2020 — expert capacity and all-to-all dispatch
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity · Fedus, Zoph & Shazeer, 2021
- ST-MoE: Designing Stable and Transferable Sparse Expert Models · Zoph et al., 2022 — router stability and auxiliary losses
- MegaBlocks: Efficient Sparse Training with Mixture-of-Experts · Gale et al., 2022 — dropless routing via block-sparse kernels
Systems behaviour, correctness and scaling
- torch.distributed — process groups, collectives and timeouts · PyTorch documentation
- PyTorch Distributed Checkpoint — sharded save and resharding on load · PyTorch documentation
- The Llama 3 Herd of Models · Grattafiori et al., 2024 — an unusually candid account of failure rates, stragglers and recovery on a 16k-GPU cluster
- An Empirical Model of Large-Batch Training · McCandlish et al., 2018 — the critical batch size
- A higher order estimate of the optimum checkpoint interval · Daly, 2006; after Young, 1974