Parallelism by the Roofline

A training step is a fight over three budgets: memory per device, compute per device, and communication between devices. Every parallelism strategy is a different way of paying those budgets for the same arithmetic.

To keep the strategies comparable, we use a single example. Most of the post comes back to one linear layer:

The cost is roughly $2 B E^2$ floating point operations. That is the compute we want to keep the hardware fed with as we spread it across devices.

Three letters carry the whole post. $B$ is the per-step batch — the number of rows of $X$, i.e. the tokens processed together. $E$ is the hidden dimension, the width shared by $X$ and $W$ and the axis the matmul contracts over. And $P$ — which shows up the moment we start splitting work across hardware — is the number of devices a given axis is spread across. A matmul’s cost grows with $B$ and with $E^2$; a collective’s cost grows with whatever it has to move. Every “useful ratio” below is just those two growth rates raced against each other.

The question every strategy answers in its own way:

What is split, what is replicated, what communication repairs the split, and can that communication be hidden under matmul compute?

The “communication” piece is always one of a small set of collectives: all-gather, reduce-scatter, all-reduce, and all-to-all. See the appendix to get an intuitive feel for these collectives.

Data Parallelism

Replicate the weights, split the batch. Each device sees a different slice of $X$‘s rows, multiplies by the same $W$, and produces its slice of $Y$:

Across chips, every $W$ is identical; each chip owns a different chunk of the batch dimension. The forward pass is literally just parallelized. A repair step comes at the end: the gradients of the replicated $W$ must take into account the full batch (which we sharded across devices), so an all-reduce of $\partial W$ fixes this.

The win is throughput, not memory. Every device still stores the full model, the full gradient, and the full optimizer state.

That backward-pass all-reduce of $\partial W$ is DP’s only cross-device cost, and — as the timeline at the end of this section shows — most of it can be hidden under compute.

The useful ratio is $B / P$. The forward pass costs nothing across devices — each chip just multiplies its own rows by its own copy of $W$. The bill comes on the backward pass: because every chip computed $\partial W$ from a different slice of the batch, the chips must agree on a single summed gradient, and that agreement is the all-reduce of $\partial W$. Here is the asymmetry that decides everything. The per-chip compute is $\sim 2 B E^2 / P$: the layer’s matmul is the $\sim 2 B E^2$ FLOPs from above, and splitting the batch across $P$ devices leaves each one with $B / P$ rows of work. But the gradient being all-reduced is “closer to $E^2$” — $\partial W$ has the same shape as $W$, an $E \times E$ grid of numbers, and that size does not depend on the batch at all. So as you add devices, the compute on each chip shrinks (it carries the $B$) while the collective does not (it carries only $E^2$). DP therefore wants $B / P$ large enough that the local matmul stays well above the collective it has to hide.

That all-reduce need not stall the backward pass. Walk two layers, $X \to Y_1 \to Y_2$: computing layer 1’s activation gradient $\partial Y_1 = \partial Y_2 \, W_2^{\top}$ needs only the local $\partial Y_2$ and the chip’s own replica of $W_2$ — not the reduced $\partial W_2$, which the optimizer only consumes at the very end of the step. So each layer’s $\partial W$ all-reduce can fire and run concurrently with the next layer’s backward compute, hiding underneath it. The timeline below shows that staircase — comm streams tucked behind compute, layer by layer. The one all-reduce that can’t hide is the first layer’s — its $\partial W$ is produced last in the backward walk, with no further backward compute to tuck it under, so a small tail stays exposed at the end of every step, and real implementations bucket gradients to keep that tail short.

FSDP

Stop replicating. The batch is still sharded as in DP, but now $W$ is sliced too — each chip stores only a strip of $W$‘s rows. In the picture below, $Y$ starts empty: chip $i$ has its $X$ slice and its $W$ shard, but it cannot compute its $Y$ yet — the matmul $X_i @ W_i$ doesn’t even compose. We must all-gather $W$ to materialize the full $W$ on every chip; only then can each chip produce its batch slice of $Y$.

The mental model that makes this work: at any moment, only one or two units are materialized (yellow in the picture above) — a unit being whatever granularity you wrap, often a layer or a small group of them. The other many units stay sharded. Memory shrinks by the shard count.

The performance question is whether the gather for layer $N{+}1$ can finish while layer $N$ is still computing. The timeline below starts in the naive schedule — each matmul waits for its all-gather to land, leaving big idle gaps in the compute stream. Click Overlapped to slide the gathers underneath the matmuls: the compute stream becomes continuous and the step takes roughly half as long.

The same per-unit all-gather will return in the backward pass — computing $\partial X$ also needs the full $W$ — so the backward pass either repeats the gather or reuses a forward-cached copy.

The useful ratio is $B / P$ — the same as DP. Per-layer compute is again $\sim 2 B E^2 / P$: the $\sim 2 B E^2$ matmul, split $P$ ways. The communication is the all-gather that rebuilds $W$ before the layer can run — and although each chip stores only its $\sim E^2 / P$ row-strip, the all-gather has to deliver the whole $E \times E$ weight to every chip, so the volume that actually moves is $\sim E^2$, the size of $W$ itself. That does not shrink as you add devices, exactly like DP’s all-reduce of $\partial W$. Race the $\sim 2 B E^2 / P$ compute against the $\sim E^2$ collective and you land on $B / P$ again. So FSDP does not win a better compute-versus-communication race than DP — if anything it pays a little more, roughly one and a half of these collectives per layer once you count the backward pass (gather $W$ going forward; gather $W$ again and reduce-scatter $\partial W$ going back). What it buys is memory: the same roofline as DP, but with the model, gradients, and optimizer state sharded $P$ ways instead of replicated on every chip. As in DP, you want the per-chip batch $B / P$ large enough that each matmul stays well above the gather it has to hide.

Tensor Parallelism

A different shape. Tensor parallelism splits the matrix itself — and there are two dual ways to slice it. Split $W$ by columns and the input stays whole; split $W$ by rows and the input must arrive already sharded. Real systems use both, in a specific order, and the pairing is the whole trick — so we take them one at a time.

The column-parallel linear is the one most readers picture. $X$ is replicated — every chip holds the full input — and each chip keeps the matching columns of $W$:

$Y_i$ is a clean slice of the output — chip $i$‘s columns, exact, no partial sums. The forward pass needs no collective at all, provided the next operation is happy consuming a column-sharded $Y$; if it needs the full thing, an all-gather rebuilds it. The bill moves to the backward pass, where the input gradient $\partial X$ must be all-reduced across chips.

The row-parallel linear is the mirror image. Now the input is sliced along the contraction dimension $E$ — each chip holds a column slice $X_i$, paired with the matching rows of $W$:

Notice the shape: $Y^{(i)}$ is the full output, not a slice. But it is a partial sum — only the $i$-th of $P$ terms in the contraction. None of the chips alone has the correct answer.

Reduce-scatter is exactly the collective the matrix shape demands: a sum across devices that re-shards the output. Above, the collective combines the striped partials and each chip is left with one column of the final $Y$.

Why carry two schemes? Because they compose. Stack a column-parallel linear in front of a row-parallel one and watch the shapes click: the first takes a replicated $X$ and emits a column-sharded intermediate — exactly the $E$-sharded input the second one wants. No collective in between; even an elementwise nonlinearity slots into the gap, since it never mixes columns. The row-parallel layer then produces partial sums, and a single all-reduce finishes the pair. Run either scheme alone and you’d pay a collective per matmul; run them as a pair — Megatron’s pattern for every MLP and attention block — and two matmuls cost one. Look back at the two diagrams: the column-parallel widget’s resting $Y$ is the row-parallel widget’s starting $X$, same shards, same colors. That is also the answer to a fair question about the row-parallel picture — “who sharded my input?”: nobody, it fell out of the previous layer that way.

In the standalone row-parallel scheme — reduce after every matmul, so the next layer starts from a clean sharded input — we cannot overlap between layers: the next layer’s input is exactly the current layer’s reduce-scattered output, so there is nothing yet to communicate ahead of. The overlap that remains happens inside the layer, with the matmul work split into a staircase that interleaves partial computation with a reduce-scatter on the previous partial. Those chunks split the token rows purely to pipeline the work — every chip still holds the full batch; only the hidden dimension $E$ is sharded across devices.

The useful ratio is $E / P$. Compute is $\sim 2 B E^2 / P$, since tensor parallelism splits one of $W$‘s $E$-sized axes across the $P$ chips — the contraction axis for row-parallel, the output axis for column-parallel. The collective — whichever scheme sends it: the all-gather of $Y$, the reduce-scatter of $Y$, or the pair’s single all-reduce — moves an activation of size $\sim B E$ — an output with one entry per (row, hidden-unit) pair, and crucially independent of $P$. Race compute against communication and the $B$ cancels, leaving a factor that scales as $E / P$: a larger hidden dimension $E$ gives the matmul more work to hide the activation exchange behind, while a larger $P$ thins each local matmul and squeezes that margin back down. The column→row pairing halves how often you pay it, but the race itself is unchanged — the ratio is $E/P$ either way.

Context Parallelism

So far the batch has been our only axis to split. But those rows are really tokens laid out along a sequence, and a long sequence is its own axis to shard. Tracking it takes on a few more dimensions: input tensors $X[B, S/P, E]$, and $Q,K,V$ tensors $[B, A, S/P, H]$ — where $A$ is the number of attention heads and $H$ the head dimension — sharded across $P$ devices. From here $B$ counts sequences rather than individual token rows: a batch element is a whole sequence of $S$ positions, so the total token count is $B \cdot S$. (You will also see this called sequence parallelism — the two terms are used interchangeably for sharding the sequence axis. Confusingly, in Megatron sequence parallelism also names a narrower trick that shards only the norm and dropout layers.) The diagrams below drop the leading batch axis $B$ — it rides along untouched and only clutters the transposes that matter — so each grid is the $[S, E]$ face of a single batch element: the sequence $S$ runs down the rows, and every device owns a horizontal band of $S/P$ of them.

For the dense layers, the split is free. A linear layer treats every token independently, so sharding the sequence is just data parallelism wearing a different hat — each device multiplies its own tokens by the replicated $W$ with no communication at all. (Free on the forward pass, at least: because $W$ is replicated, the dense layers still owe DP’s backward all-reduce of $\partial W$, exactly as in data parallelism.)

In order to compute attention, though, every query must attend to every key across the whole sequence — but each device only holds the $Q,K,V$ for its own span. That is what forces communication, and there are two different methods for repairing the split: one streams the missing keys and values past each device a block at a time, the other transposes the sharding so each device holds the whole sequence for a slice of the heads.

Ring Attention

Ring attention passes its $K,V$ shard around a ring of devices. At each step, a device attends its local $Q$ shard against whichever $K,V$ block currently sits on it and folds the result into a running output, then forwards that block onward.

The communication hides. Each device forwards its current $K,V$ block to its neighbor while it is still attending to that block, so the next hop lands just as the next matmul needs it. The first block is already local — no hop needed to start — and the last block needs no onward send, so with $P-1$ hops tucked under $P$ attention matmuls there is no exposed communication at either end — provided each hop lands inside the matmul that hides it, which is the case where the ratio below is on your side.

The useful ratio is $S/P$. Each device owns $S/P$ of the positions, so per ring hop its attention compute scales like $B \cdot A \cdot (S/P)^2 \cdot H$ — every one of its local queries attends to the $S/P$ keys in the block it holds — while the $K,V$ block it forwards is only $\sim B \cdot A \cdot (S/P) \cdot H$ numbers. Compute grows with the square of the per-device sequence length and communication only linearly, so a longer per-device span (large $S/P$) buys room to hide each hop; adding devices thins the local matmul quadratically but the hop only linearly, squeezing that margin back down. (This draws full, unmasked attention for clarity; under a causal mask the earliest blocks attend to almost nothing and the latest to everything, so production ring attention reorders the sequence — zig-zag or striped assignment — to keep every hop’s compute roughly equal.)

Ulysses Attention

Ulysses repairs the split with a different collective: an all-to-all that transposes the sharded axis. For $Q,K,V$ it turns $[B, A, S/P, H] \rightarrow [B, A/P, S, H]$ — each device trades “my positions, all heads” for “all positions, my head group”. Every device can then run full attention locally for its shard of the heads, producing $O[B, A/P, S, H]$; a second all-to-all flips it back to $O[B, A, S/P, H]$, after which the standard multi-head concatenation and output projection return $[B, S/P, E]$.

Here there is nothing to overlap. Attention’s input is exactly the first all-to-all’s output, and the next layer’s input is exactly the second’s — the same “nothing to run ahead of” bind that tensor parallelism hit between layers. As drawn, both all-to-alls are hard syncs: every device stalls through the first and again through the second. (A production kernel can chunk the transpose to claw some of that back, but the honest baseline is both all-to-alls fully exposed.)

So the roofline argument is amortization, not hiding. The exposed traffic is $\sim 4 \cdot B \cdot A \cdot (S/P) \cdot H$ (three tensors out — $Q,K,V$ — and the single output tensor back) against local attention compute of $\sim 4 \cdot B \cdot (A/P) \cdot S^2 \cdot H$ (attention is two matmuls, the $QK^{\top}$ scores and the $\cdot V$ that follows). The $P$‘s cancel and the ratio scales as $S$: you do not hide the all-to-all, you drown it under attention’s $S^2$ growth, and long sequences drown it well. The catch is structural — Ulysses gives each device one head group, so it can shard across at most as many devices as there are heads: $P \le A$ (and $P$ must divide $A$ for an even split; grouped- and multi-query attention, with their handful of $K,V$ heads, tighten the cap further). Head count, not memory, caps its reach.

Ring and Ulysses split the same job two ways, and the choice between them is the one expert parallelism will meet again. The ring hides all of its communication but ships more of it: a fresh $K,V$ block every hop, $\sim 2 \cdot B \cdot A \cdot S \cdot H$ per device once you sum the ring. Ulysses ships the bare minimum, $\sim 4 \cdot B \cdot A \cdot (S/P) \cdot H$, but pays it fully exposed. That is a factor of $\sim P/2$ more bytes for the ring, traded for hiding every one of them. Long sequences favor drowning the exposed cost (Ulysses) until you run out of heads at $P = A$; past that, only the ring keeps scaling. Both run the same collectives in reverse on the backward pass — the all-to-all sandwich inverts, the ring rotates $\partial K, \partial V$ the other way — and obey the same ratios, so we stay in the forward pass throughout.

Expert Parallelism

A Mixture-of-Experts layer replaces the single FFN with many, plus a router that sends each token to just one or two of them. The experts will not all fit on one device, so expert parallelism shards them — each device hosts a slice. (The all-to-all widget below draws the simplest case, top-1 routing — exactly one expert per token, so its combine is pure delivery. The all-gather/reduce-scatter widget draws top-2 — two experts per token, their outputs summed — because top-2 is where that method’s reduce-scatter earns the “reduce” in its name.) Just like context parallelism, there are two ways to reunite each token with its expert. One is an exact mirror of an attention method — the all-to-all sandwich is Ulysses relabeled. The other echoes ring attention rather than mirroring it: it hides its communication under compute and ships more bytes than its all-to-all sibling, the same trade the ring made, but what it moves is the tokens themselves, gathered to every device — not a $K,V$ block rotating past them.

All-to-all routing

The first method is an all-to-all sandwich, the analogue of Ulysses: dispatch every token to its expert’s device, run the expert FFN locally, then combine the outputs back home.

As drawn, both all-to-alls are exposed, for the same reason Ulysses’ are: the FFN has no input until dispatch lands, and the combine is a sync on the slowest expert — one hot expert stalls everyone behind it. The bet is amortization: each routed token pays an FFN costing $\sim E^2$ (its two $E \times E$ matmuls) while the two all-to-alls only ship its $\sim E$-sized activation there and back, so the ratio scales as $E$ — a wide hidden dimension drowns the round trip, provided the router keeps the experts balanced. Toggle the widget’s one hot expert mode to watch that proviso fail: the combine can’t complete until the overloaded device drains its queue, and everyone else waits.

All-gather / reduce-scatter routing

The second method is an all-gather / reduce-scatter sandwich, and it inverts the dispatch: instead of sending each token to its expert, send every token to every device. An all-gather replicates the token shards so each device holds the full sequence; each device then keeps just the tokens routed to its own expert, drops the rest, and runs its expert FFN — the weights never move. A reduce-scatter finishes the job, and the reduce is not decoration — the widget draws it happening. Under top-2 each token’s output is the sum of two contributions living on two different devices: each device produces a full-length output that is zero outside its own expert’s rows, and the reduce-scatter sums those summands in flight, landing each token’s finished row — one row, not two — back on its home device. (Run the same collective at the all-to-all figure’s top-1 and each sum is one real row plus zeros — routing home by arithmetic. The collective doesn’t change; the number of genuine summands does.) There is still a ring turning here, because that is how these collectives are built — shards passing neighbor-to-neighbor — but what rides it is tokens, not weights.

Now the communication hides, and it hides because we have seen both of these staircases before: the all-gather chunks and streams exactly like FSDP’s all-gather of $W$, the reduce-scatter like tensor parallelism’s reduce-scatter of $Y$, each shard’s hop tucking under the FFN chunk that overlaps it. The useful ratio is $E/P$. Per device the expert FFN costs $\sim k B S E^2 / P$ — an $\sim E^2$ FFN for each of the $\sim k B S / P$ (token, expert) pairs the router creates — while the two collectives still move $\sim 2 \cdot B S E$: the whole token set in, the whole output set back, neither of them growing with $k$. Race them and the ratio scales as $E/P$ — top-$k$ only fattens the compute side of the race — which is no coincidence: this is tensor parallelism’s ratio because this is tensor parallelism, with the experts sitting in the sharded-weight seat.

The trade against the all-to-all is bytes for exposure — the choice we met with Ulysses and the ring, in new clothes. Per device, the all-to-all ships the minimum, $\sim 2 k \cdot B S E / P$ — each token’s $k$ copies to their experts and back — and pays for it fully exposed; the all-gather/reduce-scatter ships $\sim P\times$ more, $\sim 2 \cdot B S E$, but hides every byte under compute, and its neighbor-to-neighbor traffic keeps the network calm where $P^2$ simultaneous all-to-all flows contend for it. One thing it does not buy is hot-expert immunity: each device still runs its expert over every token globally routed to it, so a skewed router overloads the hot expert’s host exactly as the all-to-all’s dispatch does — the imbalance belongs to the routing, and no collective routes around it. The gather holds one further edge: its cost is independent of top-$k$. Every token is gathered once whether it routes to one expert or four, while the all-to-all must dispatch (and combine) $k$ copies — so as $k$ grows, the $\sim P\times$ byte premium shrinks toward $\sim P/k$ and this method pulls ahead. The figure above draws exactly this: eight tokens gathered once, sixteen expert-outputs reduced back to eight rows. Backward is the transpose of forward, just as it was for FSDP and TP: all-gather the output gradients, reduce-scatter the input gradients, same collectives, same ratios.

Bringing It Together

Five strategies, one roofline. Every one of them splits the same $2 B E^2$ matmul across devices, and every one pays for the split with a collective. The whole post is the single question of whether that collective’s cost stays under the compute it rides alongside — and each “useful ratio” is just per-device compute scaling divided by collective scaling, the two growth rates raced against each other.

Two families fall out of the last two columns. The rings and staircases — DP, FSDP, TP, ring attention, all-gather/reduce-scatter experts — arrange the work so a collective always has a matmul to hide beneath, and win when the per-device compute is fat enough to cover it. The all-to-alls — Ulysses and expert dispatch — have nothing to run ahead of, so they pay the communication in the open and rely on a superlinear compute term ($S^2$ for attention, $E^2$ for the FFN) to amortize it away. Either way the knob is the same: keep the quantity that carries the compute — batch per device, hidden dimension, sequence length — large enough that arithmetic outruns the wire.

Further Reading

• JAX scaling book — training chapter (the source for much of the roofline framing here)
• PyTorch FSDP2 fully_shard
• PyTorch Tensor Parallelism
• Ring Attention with Blockwise Transformers (Liu, Zaharia, Abbeel)
• DeepSpeed-Ulysses: sequence parallelism for long-context training
• GShard: scaling giant models with conditional computation and automatic sharding
• Switch Transformers: scaling to trillion-parameter models
• NVIDIA NCCL collectives
• The Roofline Model
• Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Appendix: The Collectives

Every collective the post links to lives here, and all four now earn their place above: all-reduce, all-gather, and reduce-scatter drive DP, FSDP, and TP respectively, while the all-to-all is the workhorse of the expert-parallel routing — and the context-parallel transpose — you just stepped through.

Step through each panel; the captions describe what changes between frames.