ML Systems · Scaling the Workloads · 01 of 02

The Parallelism Playbook

There are only six things you can split across a cluster. Each split implies a parallelism strategy, and each strategy bills you for a different pattern of communication.

Distributed trainingFSDP & ZeROParallelism

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.

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.

Figure 01

One program, many ranks

Rendezvous, rank assignment, and the process groups carved out of the world afterwards. Click any GPU to see the several identities it holds at once.

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:

split samples → Data Parallelism (DDP) split model states → Sharded Data Parallelism (ZeRO / FSDP) split tensors → Tensor Parallelism (TP) split layers → Pipeline Parallelism (PP) split sequence → Sequence / Context Par. (SP / CP) split experts → Expert Parallelism (EP) Read the left column as a list of nouns in the training step — the batch, the persistent state, the weight matrices, the layer stack, the token positions, the expert set — each of which can be divided. There is no seventh entry because there is no seventh thing.

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.

Figure 02

One step, six cuts

The block is one training step: depth runs left to right, tokens run top to bottom, and the strip beneath is the persistent state each layer owns. Watch where the cut falls, and where communication has to appear as a result.

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:

  1. What is replicated? Whatever the cut left whole on every device.
  2. What is sharded? Whatever the cut divided.
  3. What memory does it save? The sharded part divided by the degree, and nothing else.
  4. What computation remains? Whether each device still does the full arithmetic or only a fraction of it.
  5. What communication appears? The collective needed to repair the dependency the cut severed.
  6. How frequently? Once per step, once per layer, or once per operation. This matters more than volume.
  7. What topology does it prefer? Frequent and blocking wants a scale-up fabric; infrequent and overlappable tolerates scale-out.
  8. 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.

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:

T_message ≈ α + n / β α is fixed per-message latency, n is bytes, β is bandwidth. Small messages are α-bound and want fusing; large ones are β-bound and want fast links. Ring algorithms make per-rank traffic converge to a constant multiple of the buffer size regardless of world size — so what actually grows with N is the latency term, and the chance that one rank is slow.
Figure 03

The collectives, as shapes

A refresher rather than a derivation. Each coloured square is a chunk, each row a rank. The movement is the definition — and the shape of the movement, not its volume, is what a network topology punishes.

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:

L(θ) = (1/B) Σ_{j=1..B} ℓ(x_j; θ) = (1/N) Σ_{i=1..N} [ (1/(B/N)) Σ_{j ∈ shard i} ℓ(x_j; θ) ] = (1/N) Σ_i L_i(θ) ∇L = (1/N) Σ_i ∇L_i The gradient of the global batch is exactly the average of the per-rank gradients. So the operation we need is: combine one tensor elementwise across all ranks, with the result available everywhere, before anyone steps. That is the definition of an AllReduce. We did not choose it — the arithmetic did.

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:

global batch = microbatch_per_device × DP_degree × accumulation_steps The first is bounded by device memory. The second by how many GPUs you have. The third only by patience — accumulation buys effective batch size in time rather than hardware, and with 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.

Figure 04

Where a global batch comes from — and how a sampler breaks it

The dataset is partitioned by rank, consumed in microbatches, and accumulated into one optimizer step. Switch the sampler to a broken variant and watch duplicates appear while every other number stays reassuringly correct.

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.

Figure 05

Gradient buckets firing during the backward pass

Two lanes: compute above, communication below. Change the bucket size and watch the trade — few large buckets get better bandwidth but start later; many small ones start early but pay α repeatedly and can outrun the link.

A few consequences worth knowing:

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.

Figure 06

The three resources a "free" overlap actually consumes

Nominal overlap against achieved overlap, with SM contention and link sharing turned on one at a time. The gap between the two bars is the difference between a timeline diagram and a profiler trace.

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.

Figure 07

Adding GPUs does not add memory

Left: throughput rising with world size. Right: per-device memory, flat as a board. The whole of the next section is an attack on that flat line.

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.

comm / compute ∝ Ψ / (b_local · s · Ψ) = 1 / (b_local · s) More tokens per device means better overlap. This is why "increase the per-device batch" is the first and most effective fix for a communication-bound DDP run — and why it stops being available exactly when memory gets tight. The two constraints are coupled, which is the recurring theme of the whole article.

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 Ψ:

per-device bytes comm volume DDP / ZeRO-0 2Ψ + 2Ψ + 12Ψ = 16Ψ 2Ψ (AllReduce) ZeRO-1 2Ψ + 2Ψ + 12Ψ/N = 4Ψ + 12Ψ/N 2Ψ (RS + AG) ZeRO-2 2Ψ + (2Ψ + 12Ψ)/N = 2Ψ + 14Ψ/N 2Ψ (RS + AG) ZeRO-3 (2Ψ + 2Ψ + 12Ψ)/N = 16Ψ/N 3Ψ (2×AG + RS) Stage 1 shards optimizer state, which is only touched during the update — so each rank simply updates its own slice. Stage 2 additionally shards gradients: a rank only needs the gradient slice matching its optimizer slice, and that slice is exactly what a ReduceScatter produces. Stage 3 shards parameters, which is the first stage that changes the forward pass, because a rank can no longer compute a layer without first fetching its weights.

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.

Figure 08

The ZeRO ladder, one shard at a time

Watch the per-device ledger collapse as each stage sheds a term, and the communication counter climb when stage 3 turns on. Set the shard degree, the parameter count and the optimizer, and the numbers follow.

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:

forward: parameter shards ↓ AllGather full layer → compute → reshard (free the full tensor) backward: AllGather parameters again ↓ compute gradient → ReduceScatter → keep only my slice step: update my own shard. No communication at all. Outside computation, each rank holds only shards. Memory becomes a streaming working set rather than a static allocation, and peak usage is set by the largest unit's gathered size plus its activations — not by the model.

Several policy knobs follow directly from this rhythm, and each has a characteristic failure:

Figure 09

FSDP with and without prefetch

The same layer sequence run twice. Above: gather, compute, free, strictly in order. Below: the gather for layer i+1 issued during layer i. Same bytes, very different timeline.

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.

Figure 10

HSDP: putting the chatty collective on the fast wire

Adjust the shard group size and the two fabrics' bandwidths. The optimum is almost always "shard exactly as wide as you must to fit, and no wider".

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

Y = X A X: [b·s, h] A: [h, 4h] split A by columns: A = [A₁ | A₂] Y = X [A₁ | A₂] = [X A₁ | X A₂] = [Y₁ | Y₂] Rank 1 computes Y₁ = X A₁, rank 2 computes Y₂ = X A₂. Both need the full X — which they already have, because X is the block's replicated input. No communication in the forward pass at all. The output is sharded along its feature dimension.

Row parallel: split the input features

Z = Y B Y: [b·s, 4h] B: [4h, h] split B by rows and Y by columns: Z = [Y₁ | Y₂] · [B₁ ; B₂] = Y₁B₁ + Y₂B₂ Each rank computes a partial sum of the full-shaped output, and those pieces must be added: one AllReduce. Crucially, row-parallel consumes a column-sharded input — exactly what the previous layer produced. The two forms are designed to chain.

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.

Figure 11

Column-parallel then row-parallel, block by block

The actual matrix blocks, moving. Follow one activation row through the MLP and watch where the partial sums appear — the AllReduce is not a design choice, it is the only place the arithmetic can be completed.
per transformer layer, per step: forward: AllReduce after attention + AllReduce after MLP = 2 backward: AllReduce for each block's input gradient = 2 message size: b·s·h — the activation, not the parameter count Megatron writes the block boundaries as conjugate operators f and f̄: f is identity forward / AllReduce backward, f̄ is the reverse. Because the message is activation-sized, tensor-parallel traffic grows when you make the step bigger — the opposite of DDP, where it is fixed by the parameter count.

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.

Figure 12

The tensor-parallel tax, on the fast fabric and off it

Per-layer exposed communication as the TP degree crosses the node boundary. The discontinuity is not a modelling artefact — it is the fabric.

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.

TP region : activations sharded by hidden dim, sequence complete SP region : activations sharded by sequence, hidden dim complete TP → SP boundary : ReduceScatter (instead of AllReduce) SP → TP boundary : AllGather (instead of nothing) Here is the elegant part: an AllReduce is a ReduceScatter followed by an AllGather. Splitting it in half and putting the sequence-parallel region between the halves moves the same bytes — the collective count and volume are unchanged. You get the memory saving for free, which is why sequence parallelism is essentially always on when tensor parallelism is.

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.

Figure 13

Alternating layouts: TP regions and SP regions

One transformer block, following the activation tensor through both layouts. Toggle sequence parallelism to see which tensors stop being replicated — and confirm that the bytes on the wire do not change.

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.

online softmax accumulation, block by block: m_new = max(m_old, rowmax(S_k)) O_new = O_old · exp(m_old − m_new) + exp(S_k − m_new) · V_k ℓ_new = ℓ_old · exp(m_old − m_new) + rowsum(exp(S_k − m_new)) This is why ring attention is possible at all: softmax over a partitioned axis can be computed incrementally with a running maximum and a running denominator, so no rank ever needs a full score row. The result is exact, not approximate — the same identity FlashAttention already uses, applied across devices instead of across tiles.

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.

Figure 14

Ring attention and Ulysses, side by side

Switch between schemes on the same sequence. Ring rotates KV blocks around the circle; Ulysses transposes the layout so attention becomes local. Watch the causal triangle to see why naïve ring assignment leaves rank 0 idle.

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.

GPipe schedule: all m forwards, then all m backwards busy time = m · (t_f + t_b) bubble time = (P−1) · (t_f + t_b) bubble ratio = (P−1) / (m + P − 1) With P=8 stages and m=8 microbatches you waste 47% of the cluster. With m=64 you waste 10%. The rule of thumb m ≥ 4P comes straight from this expression — and it is why pipeline parallelism wants large global batches, which means it competes with data parallelism for the same batch budget.

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.

Figure 15

Four pipeline schedules on one clock

GPipe, 1F1B, interleaved and zero-bubble drawn to the same clock. Change P and m and watch the bubble ratio move; the memory readout is the number of live activation stashes on the worst-off stage.

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.

router: g = softmax(x · W_r) W_r: [h, E] pick top-k experts per token y = Σ_{e ∈ topk} g_e · Expert_e(x) capacity per expert = capacity_factor · k · T / E If routing were uniform each expert would get exactly k·T/E tokens. It is not, so implementations allocate headroom via a capacity factor and drop tokens that overflow — they skip the expert and pass through on the residual. The capacity factor trades wasted compute on under-filled experts against information lost to dropped tokens.

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.

Figure 16

Token routing, capacity, and the two all-to-alls

Tokens flow to experts, overflow gets dropped at the capacity line, results flow back. Turn the balancing loss down and watch the distribution collapse — and the wasted capacity climb with it.

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.

N = DP × TP × PP × CP × EP 512 = 8 × 8 × 8 × 1 × 1 for example rank → coordinate, with TP fastest-varying: tp = r % TP cp = (r / TP) % CP dp = (r / (TP·CP)) % DP pp = r / (TP·CP·DP) The ordering is the whole design. Ranks adjacent in the linear order are physically adjacent — same node, same switch. So whichever axis you place fastest-varying gets the fastest links. Every axis also owns a distinct process group, and using the right collective on the wrong group is a category of bug that exists only in distributed code.
Figure 17

The device mesh, explorable

Set the five degrees and watch the ranks fold into a mesh. Selecting a rank lights up its communication groups, coloured by the fabric each one actually rides. Configurations that place a TP group across a node boundary are flagged.

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:

  1. 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.
  2. 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.
  3. 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.
  4. 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.
  5. 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:

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".

Figure 18

Where each strategy's communication sits

Every axis on the same step timeline, positioned by how far its collective sits from the operation that consumes it. Distance from the consumer is the entire story of whether a collective is free.

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.

Figure 19

The decision tree

Answer at each branch and watch the path light up. The tree is deliberately shallow: most real configurations are two or three nodes deep, and the ones that go deeper are paying for it.

Written as a ladder, with the step everyone skips at the top:

  1. 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.
  2. If the model fits, use data parallelism — with ZeRO-1, which is free.
  3. If model states do not fit, shard them: FSDP/ZeRO-3, at the smallest degree that fits, replicating beyond it.
  4. If one layer is itself too large, or you need more compute per layer, add tensor parallelism — inside the fast-link domain only.
  5. If depth still exceeds memory, or you need model partitioning across nodes, add pipeline parallelism.
  6. If long context dominates activation memory, add context parallelism — after activation checkpointing, which costs predictable compute rather than unpredictable bandwidth.
  7. If the architecture is MoE, add expert parallelism and design the topology around All-to-All.
Communication volumes are per device per step; Ψ is the parameter count, and activation-sized traffic scales with b·s·h.
StrategyWhat is splitMemory effectCommunicationSync granularityReach for it when
DDPthe batchnone — 16Ψ replicatedAllReduce, ~2× grad bytesbackward bucketsthe model already fits
ZeRO-1/2optimizer, gradients16Ψ → 4Ψ, then 2Ψsame as DDPbackward bucketsalways — it is free
FSDP / ZeRO-3everything persistent16Ψ/N1.5× DDP, per layerper sharded unitmodel states do not fit
Tensoreach matmulweights and inner activations / T4 AllReduces per layer, activation-sizedsub-layer, blockingindividual layers are huge; stays in-node
Pipelinethe layer stackweights / PP2P, one activation per boundarypipeline clockdeep models; slow network
Contextthe sequenceactivations / CKV ring or 4 all-to-allsattention layeractivation OOM from long sequences
Expertthe expertsexpert weights / E2 All-to-Alls per layerMoE routingmixture-of-experts models
Figure 20

The same seven strategies, as one moving picture

A guided tour of the table above: each strategy as what it splits, what it costs, and where it sits on the memory/communication plane.

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.

Figure 21

The planner lab

A first-order model. It will get you to a sensible starting configuration and will not replace a profiler — the interconnect assumptions alone can be off by a factor of two on a busy cluster.

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:

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.

Figure 22

Checkpoint geometry and resharding

One optimizer tensor, sharded under one mesh and reloaded under another. Compare gathered against sharded writes, then change the world size on restart and watch which layouts survive.

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:

η_N = T₁ / (N · T_N) T₁ is the time for the work on one device, T_N the time on N. Perfect scaling gives η = 1. The definition is only meaningful once you say which work is held constant — which is the difference between strong and weak scaling, and the reason the two curves diverge so sharply.

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.

Figure 23

Strong scaling, weak scaling, and the critical batch size

Three curves usually plotted separately that should not be. Push the world size past the critical batch marker and watch tokens-per-second keep climbing while tokens-to-target-loss climbs with it.

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".

Figure 24

How a distributed run hangs

Four ranks on one timeline. Pick a failure mode and watch the collective fail to close — three ranks parked inside an AllReduce that can never complete, and a watchdog counting down.

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.

T_step = max_i (t_i) over N ranks If each t_i has mean μ and a modest tail, E[max] grows roughly like μ + σ·√(2 ln N). At N=8 the penalty is negligible. At N=4096 it is substantial even with a tight distribution. This is why large runs invest so heavily in per-rank timing telemetry: finding the one bad card is often worth more than any algorithmic optimisation.
Figure 25

The straggler tax

Per-rank step times drawn from a distribution, with the cluster's step time set by the slowest. Increase the world size without changing the distribution at all, and watch the tax grow anyway.

The rest of the list

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.

Tensor, sequence and pipeline parallelism

Mixture of experts

Systems behaviour, correctness and scaling