Part V — The hard parts

Distributed compilation

Lesson 12 handed you the joint forward+backward graph and ran the whole middle end over it — fused, planned, laid out, scheduled. Every one of those passes optimized for one device. But the models that need this machinery are exactly the ones that don't fit on one device: a 70B's training state is ~1.1 TB, a single MLP block's transient footprint can overflow 80 GB, and even when it fits, one GPU finishes a step too slowly to ever train. The classic answer is to shard by hand — split this matmul across eight GPUs, chain those layers into a pipeline, all-reduce the gradients — and the result is some of the most error-prone code in the field. This lesson makes partitioning the compiler's job: you annotate a handful of tensors, and the compiler propagates the sharding through the graph, inserts the collectives, and overlaps them with compute.

The previous step left this broken

The joint graph from lesson 12 is a single-device program: it assumes every tensor lives in one HBM and every op reads its inputs locally. For a frontier model that assumption is simply false — the parameters, the optimizer state, and the activations each exceed one GPU. The hand-written fix (cross-link cs336 · 08 data parallelism and cs336 · 09 model parallelism) demands that a human reason globally about where every tensor is split, insert the right collective at every layout mismatch, get the backward's collectives right by symmetry, and place each axis on the link it can afford — across hundreds of devices, by hand, with a single wrong all-reduce silently corrupting the gradient. It does not compose with the compiler: the moment you fuse or repartition a layer, the hand-placed collectives are wrong. Partitioning by hand is intractable at scale and fights every pass you just built.

Linear position

Forced by: the joint graph (lesson 12) exceeds one device, and placing tensors and inserting collectives across hundreds of GPUs by hand is intractable and doesn't compose with the rest of the compiler.
New idea: make partitioning a compiler pass — annotate a few tensors with sharding specs over a device mesh, then the compiler propagates shardings through the graph (GSPMD / Shardy / jax.shard_map / DTensor), inserts the collectives (all-reduce, all-gather, reduce-scatter, all-to-all) wherever the producer's and consumer's layouts disagree, and schedules comm/compute overlap as just another scheduling pass (call back lesson 08). Single-device semantics + sharding hints → an SPMD program for N devices.
Forces next: you've now seen every pass that makes a model fast and big — capture, IR, rewrite, fuse, plan, layout, schedule, codegen, autotune, dynamic shapes, autodiff, and now distribution. But every one of them quietly assumed the compiler can generate every op and that emitting kernels is the finish line. Shipping breaks that assumption, and the first crack is the long tail of ops the compiler can't generate (lesson 14).

The plan

Five moves. (1) Why data/tensor/pipeline parallelism by hand is brittle — and the alternative: declare the sharding and let the compiler do the rest. (2) The vocabulary — device mesh, sharding spec (replicate / shard along an axis), and SPMD. (3) Sharding propagation: given a few annotations, infer the rest by minimizing communication, and insert a collective at every mismatch — with worked comm-byte numbers for a sharded matmul. (4) Collective scheduling: overlap comm with compute (it's a scheduling pass) and fuse collectives. (5) The stacks — XLA GSPMD/Shardy, PyTorch DTensor + torch.compile, Megatron as the hand-written contrast. Then drive the propagator until a bad annotation needs an all-to-all everywhere.

1 · The problem: hand-sharding is brittle and doesn't compose

The three parallelism axes from cs336 are real and necessary. Data parallelism replicates the model and splits the batch; tensor parallelism splits each matmul (column-parallel then row-parallel, one all-reduce per block); pipeline parallelism splits the layer stack into stages. Written by hand — the Megatron way — each is a separate, intricate rewrite of the model code: you replace nn.Linear with ColumnParallelLinear, you hand-place all_reduce calls, you write the backward's collectives to mirror the forward, and you wire the micro-batch pipeline schedule. It works, and for the configurations Megatron ships it is close to optimal. But it is rigid: change the mesh shape, fuse two layers, add an expert, and the hand-placed collectives are now in the wrong places.

The compiler's bet is different. You wrote a single-device program — one tensor, one matmul, one HBM — and that program already states the math exactly. The only new information distribution needs is where each tensor is split. So: keep the single-device program, attach a sharding annotation to a handful of tensors (the inputs, the parameters), and let the compiler derive everything else. Partitioning becomes a pass in the same pipeline as fusion and memory planning, subject to the same invariant — it must preserve the math exactly (the motif from lessons 03–04) — which is precisely why inserting a collective is legal: an all-reduce that sums partial products computes the same value the single-device matmul would.

2 · The vocabulary: device mesh, sharding spec, SPMD

A device mesh is your N devices arranged as a named, multi-dimensional grid. Sixteen GPUs might be a 1-D mesh (16,), or a 2-D mesh (4, 4) with axes named data and model. The axis names are the handle you shard against — "split this tensor's rows across the model axis" — and they let one annotation scale to any device count.

A sharding spec says, for each dimension of a tensor, whether it is replicated (every device holds the full copy) or sharded along a mesh axis (each device holds a contiguous slice). Write a matrix's sharding as a tuple over its dims: (model, replicate) means dim 0 is split across the model axis and dim 1 is whole on every device. A 8,192 × 8,192 bf16 weight (128 MB) sharded (model, replicate) over a model axis of 8 lives as eight 1,024 × 8,192 slices of 16 MB each — that is how the parameters fit.

device mesh

N devices as a named grid, e.g. (data=4, model=4). Axis names are what you shard against, so one spec scales to any size.

sharding spec

per-tensor-dim: replicate (full copy everywhere) or shard along a mesh axis (a slice per device). The compiler's currency.

SPMD

single-program-multiple-data: one program text runs on every device, each on its own shard. The compiler emits it from single-device code.

collective

a comm op every device runs together: all-reduce, all-gather, reduce-scatter, all-to-all. Inserted wherever two layouts disagree.

SPMD — single-program-multiple-data — is the execution model that makes this affordable. Rather than emit a different program per device, the compiler emits one program that every device runs, each on its own shard, with collectives at the synchronization points. The host launches the same compiled artifact N times; the only difference between devices is which slice of data they hold. This is what lets a 1,000-line model compile once and run on 8 or 8,000 GPUs.

3 · Sharding propagation: infer the rest, insert collectives at mismatches

You annotate a few tensors; the compiler must decide the sharding of every other tensor in the graph. This is sharding propagation (GSPMD's core algorithm): treat shardings as a dataflow property and push them forward and backward through the ops until every value has a spec, choosing, at each op, the layout that minimizes communication. It is the same shape-inference idea as the type system in lesson 03, but the property being inferred is "where is this split," and the cost being minimized is bytes-on-the-wire.

The crux is a matmul whose operands disagree. Consider Y = X · A with the contraction dimension K. If X is sharded along K and A is sharded along the same K, then each device computes a partial product over its slice of K, and the true result is the sum of those partials → the compiler inserts one all-reduce. This is exactly Megatron's row-parallel pattern, derived automatically. But annotate badly — say X sharded on K while A is sharded on its output dim N — and the layouts don't line up for a local matmul at all; the compiler must first reshard one operand with an all-to-all (every device sends a piece to every other), which moves far more bytes, before any math runs. The annotation is the communication bill.

Worked comm bytes. Take X: 8,192 × 8,192, A: 8,192 × 8,192, bf16, on a model axis of 8 devices. The output Y is 8,192 × 8,192 × 2 ≈ 134 MB.

annotation	what the compiler must insert	comm bytes moved
good: both operands sharded on contraction K	one all-reduce of the 134 MB partial result	≈ 2 · 134 MB ≈ 268 MB (ring all-reduce moves ≈2× the tensor)
good: weights sharded on output N, X replicated	nothing for the matmul; an all-gather only if a later op wants Y whole	0–134 MB, deferred
bad: operands misaligned (X on K, A on N)	all-to-all to reshard, then the matmul, then maybe another collective	≈ 134 MB reshuffled + the matmul's own collective — multiplied at every layer

The lesson of the table: a sharded matmul's cost is set by the annotation, not by the FLOPs. A good sharding closes a two-matmul block (attention or MLP) with a single all-reduce — the Megatron cost — and the propagator finds it for you. A bad one pays an all-to-all reshard at every layer, and across 80 layers that is the difference between a step that overlaps its comm and one that is entirely bandwidth-bound. The collectives, named once:

all-reducesum (or max/min) partials across devices, every device gets the full result. The row-parallel matmul's closer; ≈2× tensor bytes on a ring.

all-gatherconcatenate every device's shard so all end up with the whole tensor. Used to un-shard a value a later op needs replicated.

reduce-scattersum partials and leave each device only its slice — half of an all-reduce; the workhorse of FSDP/ZeRO gradient sharding.

all-to-allevery device sends a distinct piece to every other — a full transpose of the sharding. The expensive reshard; also MoE expert routing.

4 · Collective scheduling — overlap comm with compute

Inserting the right collectives makes the program correct; it does not make it fast. A naively placed all-reduce blocks the next op on the critical path, and the device idles while bytes cross the wire. The fix is the same one lesson 08 used for loop nests: treat collective placement as a scheduling problem. Overlap means running a collective concurrently with independent compute on a separate stream, so the comm latency is hidden under math the device was going to do anyway. The data-parallel gradient all-reduce is the textbook case — it overlaps almost perfectly with the backward pass, because gradients for early layers are ready while later layers are still computing, so the all-reduce of layer L's grads runs while layer L−1's grads are still being produced.

Two scheduling moves the compiler makes here. First, reorder for overlap: hoist a collective earlier so more independent compute sits between its launch and its first use, maximizing the hidden window. Second, collective fusion: merge many small collectives into one large one — a hundred tiny per-tensor gradient all-reduces become a few bucketed all-reduces, because a collective's cost is latency + bytes/bandwidth, and batching amortizes the fixed latency. (This is the comm analogue of operator fusion from lesson 05: fewer, larger operations beat many small ones for the same total bytes.) When the compiler owns partitioning, it owns this schedule too — which is exactly why hand-sharding leaves performance on the table: a human places the collective correctly but rarely overlaps it as aggressively as a pass that can see the whole dependency graph.

5 · The stacks: GSPMD, DTensor, and Megatron as the contrast

Three points on the spectrum from "compiler does everything" to "human does everything":

system	how you express sharding	who inserts collectives	niche
XLA GSPMD / Shardy	annotate a few tensors with sharding specs over a mesh (`jax.sharding`, `jax.shard_map`); the compiler does the rest	the compiler — full propagation + collective insertion + overlap	JAX / TF on TPU & GPU; the most automatic. Shardy is the next-gen MLIR-based propagation system replacing GSPMD.
PyTorch DTensor + torch.compile	wrap tensors as `DTensor` with a placement (`Shard(dim)` / `Replicate()`) on a `DeviceMesh`; eager redistributes, `torch.compile` fuses/overlaps	DTensor inserts collectives at redistribution; Inductor schedules overlap	PyTorch-native; composes with FSDP2 and the rest of the eager ecosystem
Megatron-LM	you rewrite the model with `ColumnParallelLinear` / `RowParallelLinear` and hand-placed collectives	the human — collectives are written into the model code	peak performance for its fixed recipes; the hand-written baseline the compilers aim to match

The verdict the field has reached: the compiler approach (GSPMD/Shardy, DTensor) wins on flexibility and composition — you change one annotation to repartition, and the sharding pass re-runs against fused, autodiffed, dynamically-shaped graphs without you rewriting collectives. Megatron still wins on peak for the exact configurations it was tuned for, because a human hand-picked every collective and overlap. The trajectory is toward the compiler: as propagation and scheduling get better, the hand-written gap shrinks, and the flexibility of "annotate, don't rewrite" matters more as meshes and models change every few months.

Drive it: the sharding propagator

The widget is a two-matmul graph — H = X · A then Y = H · B, the shape of an MLP block — on a device mesh you size to 2, 4, or 8. You pick how the input X is sharded; the propagator infers the rest, inserts the collectives at the mismatches, and reports the bill. Set a good annotation (operands aligned on the contraction dim) and the block closes with one all-reduce, comm bytes stay flat as the mesh grows, and per-device memory drops by the mesh size. Set a bad one (misaligned) and watch all-to-all reshards appear between every op, comm bytes explode, and the overlap fraction collapse. The knob that breaks it: flip "shard X on the wrong axis" and the green all-reduce turns into red all-to-alls everywhere.

Sharding propagator — annotate one tensor, pay the comm bill

A two-matmul block X·A·B on an N-device mesh. Pick the input sharding; the compiler propagates it, inserts collectives at every layout mismatch, and prices them. A good annotation closes the block with one all-reduce (flat comm, memory ÷ N). A bad one forces an all-to-all reshard at every op — comm explodes and overlap dies. Flip the annotation and watch the graph go red.

mesh size N: input X sharding: tensor dim d: 8,192

collectives inserted

—

comm bytes / step

—

per-device memory

—

compute / comm overlap

—

layout OK

The shape of the readouts is the whole point: the same math, the same FLOPs, the same mesh — and the only thing you changed was a single sharding annotation, yet the comm bill moved by an order of magnitude. That is why partitioning belongs in the compiler: the propagator searches for the layout that minimizes communication so you don't have to hand-place a single collective.

Failure modes & checklist

Failure modes

Annotations that fight the math. Sharding operands on mismatched axes so the propagator must reshard with all-to-all at every op. Signal: the trace is dominated by all-to-all collectives, not matmuls; step time barely improves as you add devices.
Trusting propagation blindly on a huge graph. Under-annotating so the compiler picks a globally poor layout it can't escape. Signal: unexpected all-gathers of large activations; comm bytes far above the Megatron one-all-reduce-per-block baseline.
No overlap. Collectives placed correctly but on the critical path, so the device idles on the wire. Signal: a gap in the timeline exactly the length of each all-reduce; MFU well below what FLOPs predict.
Tensor-parallel collectives crossing a node boundary. Mapping the heavy all-reduce axis onto the slow inter-node link. Signal: per-step time scales with inter-node bandwidth, not NVLink (cross-link cs336/09 — TP ≤ GPUs/node).
Mismatched forward/backward shardings. Letting autodiff (lesson 12) produce a backward whose layouts disagree with the forward, forcing extra reshards. Signal: the backward has more collectives than the forward for a symmetric block.

Checklist

Annotate the few load-bearing tensors (inputs, parameters) and let propagation infer the rest — don't hand-spec every value.
Align operands on the contraction dim so a two-matmul block closes with one all-reduce, the Megatron-optimal cost.
Place the heavy axis on the fast link — tensor-parallel collectives on NVLink intra-node, data-parallel reduce-scatter on the wide cheap axis.
Confirm overlap in the profile: every collective should hide under independent compute; bucket small collectives via fusion.
Check the comm bill against the baseline — one all-reduce per attention/MLP block forward and backward; more means the propagation chose a poor layout.

Checkpoint

Try it

Open the widget at mesh size 4, input sharding "good", d = 8,192. Read the three KPIs: one all-reduce, comm bytes around a few hundred MB, per-device memory cut to a quarter, overlap high. Now switch input sharding to "bad" and watch all-to-all collectives appear between every op and comm bytes jump roughly an order of magnitude while overlap collapses. Reset to "good" and slide the mesh from 2 → 8: per-device memory keeps dropping (the parameters shard further) while comm bytes stay roughly flat — that's the win. By hand, for the good case: a 8,192 × 8,192 bf16 result is ≈134 MB, and a ring all-reduce moves ≈2× the tensor, so predict the comm bytes and check the KPI. Finally explain in one sentence why the bad annotation's cost is multiplied by the layer count while the good one is not.

Where this points next

You can now take a single-device graph and, from a handful of sharding hints, get a correct, overlapped SPMD program for hundreds of devices — partitioning is a pass like any other, and it composes with fusion, autodiff, and dynamic shapes because they all share the preserve-the-math invariant. That completes the catalog: you have seen every pass a modern ML compiler runs, each one forced by the wall the previous one hit. But all of those passes share an unspoken assumption: that the compiler can generate every op in the graph, and that once kernels are emitted the job is done. That is exactly what shipping a model to production breaks. The next lesson opens Part VI — making the compiled model shippable — and starts with the first crack: the long tail of ops the compiler cannot generate, and what to do when it hits one. That is 14 · Operator coverage — decomposition & the long tail.

Takeaway

Lesson 12's joint graph is optimized for one device, but frontier models exceed any single GPU and hand-sharding them — splitting matmuls, hand-placing collectives, mirroring the backward, across hundreds of devices — is intractable and doesn't compose with the rest of the compiler. The fix is to make partitioning a compiler pass: annotate a few tensors with sharding specs (replicate or shard along a named axis) over a device mesh, and the compiler propagates shardings through the graph (GSPMD / Shardy / DTensor), inserts a collective — all-reduce, all-gather, reduce-scatter, all-to-all — at every layout mismatch, and turns single-device semantics into an SPMD program every device runs on its own shard. The annotation is the comm bill: operands aligned on the contraction dim close a two-matmul block with one all-reduce of ≈2× the result tensor (Megatron's cost, derived automatically), while a misaligned annotation forces an all-to-all reshard at every layer. Then collective scheduling — the same scheduling idea as lesson 08 — overlaps comm with independent compute and fuses small collectives into bucketed ones, hiding the wire under math. XLA GSPMD/Shardy and PyTorch DTensor automate it; Megatron is the hand-written contrast that still wins on peak for fixed recipes. The compiler turns "single-device program + sharding hints" into N-device SPMD — and with that you have seen every pass that makes a model fast and big, each one quietly assuming the compiler can generate every op. Lesson 14 opens Part VI and breaks that assumption: making the model shippable starts with the long tail of ops the compiler cannot generate.

Interview prompts

What is a sharding spec and a device mesh, and why annotate only a few tensors? (§2–3 — a mesh is N devices as a named grid; a spec says, per tensor dim, replicate or shard along a mesh axis. You annotate a few load-bearing tensors and the propagator infers the rest by minimizing communication, so one spec scales to any device count.)
How does sharding propagation decide where to put a collective? (§3 — it pushes shardings through the graph as a dataflow property; at each op it picks the layout minimizing comm, and inserts a collective wherever producer and consumer layouts disagree — e.g. operands sharded on the contraction dim need one all-reduce to sum partials.)
Why does a good vs bad annotation change the comm bill by an order of magnitude? (§3 — aligned operands close a block with one all-reduce (≈2× the result tensor); misaligned operands force an all-to-all reshard before the matmul, moving far more bytes, and that cost repeats at every layer.)
Distinguish all-reduce, all-gather, reduce-scatter, all-to-all and where each appears. (§3 — all-reduce sums partials so all get the full result (row-parallel matmul); all-gather concatenates shards to replicate; reduce-scatter sums and keeps a slice each (FSDP grads); all-to-all is a full sharding transpose (resharding, MoE routing).)
Why is collective scheduling the same problem as loop scheduling? (§4 — placement decides performance: a collective on the critical path idles the device, so you reorder to overlap it with independent compute on a separate stream and fuse small collectives into bucketed ones — exactly the scheduling/fusion moves from lessons 08 and 05, applied to comm.)
Compare GSPMD/Shardy, DTensor, and Megatron. (§5 — GSPMD/Shardy: annotate, compiler propagates + inserts + overlaps, most automatic; DTensor: placements on a mesh, collectives at redistribution, torch.compile overlaps; Megatron: human rewrites the model with parallel layers and hand-placed collectives — peak for fixed recipes, the baseline compilers chase.)
Why does compiler-based partitioning compose with the rest of the compiler when hand-sharding doesn't? (§1 — inserting a collective preserves the math, so it's just another semantics-preserving pass; it re-runs against fused, autodiffed, dynamically-shaped graphs, whereas hand-placed collectives break the moment you repartition or fuse a layer.)