Part II — Systems

Model parallelism — tensor, pipeline, the 3D mesh

Lesson 08 sharded the storage: ZeRO-3 / FSDP cut the 16-bytes/param training state by the data-parallel degree, so a 7B that overflowed one 80 GB GPU now fits across eight. But every FSDP rank still executes the whole model on its own slice of the batch — it gathers each layer's params just-in-time, runs the full forward, then frees them. That breaks the moment a single layer is too big to compute on one device, the activations for one layer overflow, or the model is simply larger than any GPU even after gathering. Data parallelism can shrink where weights live; it can never split where the math runs. This lesson splits the model itself.

1 · Tensor parallelism — split the matmul itself

The unit of work inside a Transformer is a matmul: Y = X·A with A a weight matrix. Tensor parallelism (TP), introduced by Megatron-LM, splits A across t GPUs so each holds a slice, computes a partial result, and the partials are combined by a collective. The art is arranging the cuts so that two matmuls in a row — which every Transformer block has — cost exactly one communication instead of two.

Column-parallel then row-parallel. Take the MLP, two matmuls back to back: Z = GELU(X·A)·B. Split the first matrix A by its columns into [A₁ A₂ … A_t]; GPU i computes GELU(X·A_i) — a full slice of the hidden activation, no communication needed because GELU is elementwise. Now split the second matrix B by its rows into [B₁; B₂; …; B_t] so that GPU i already holds the matching activation slice and computes the partial product GELU(X·A_i)·B_i. The full output is the sum of those partials — one all-reduce across the t GPUs, and you're done. Two matmuls, one collective.

Attention takes the identical pattern. The QKV projection is column-parallel: assign whole attention heads to GPUs, so GPU i owns a subset of heads and computes their Q, K, V and the full per-head attention locally — heads are independent, so no cross-GPU comm inside the softmax. The output projection W_o is then row-parallel, and a single all-reduce sums the per-head contributions back into the residual stream. So the canonical Megatron block costs two all-reduces in the forward (one for attention, one for the MLP) and, by symmetry, two in the backward — four collectives per layer, each over the t TP ranks.

  MLP block, t=2 tensor-parallel GPUs
                 X  (replicated on both)
        ┌──────────────┴──────────────┐
   GPU0: X·A₁ (cols 0..h/2)       GPU1: X·A₂ (cols h/2..h)
        GELU (elementwise, local)      GELU (local)
   GPU0: (·)·B₁ (rows 0..h/2)     GPU1: (·)·B₂ (rows h/2..h)
        └──────────────┬──────────────┘
              ALL-REDUCE (sum partials)  ← the one collective
                        Z  (replicated)

Why TP must stay inside a node. That all-reduce is not cheap and it sits on the critical path — the matmul cannot finish until the bytes arrive. Each forward all-reduce of the residual stream moves on the order of B·T·d elements; with four collectives per layer over dozens of layers, TP communicates an enormous amount per step, and it is synchronous (no useful overlap — the next op depends on the result). The only way this is affordable is on an intra-node fabric: NVLink/NVSwitch gives ~900 GB/s per GPU between the 8 GPUs in a node, versus ~25–50 GB/s per GPU over InfiniBand between nodes — roughly a 20–40× gap. Span TP across a node boundary and every block stalls on the slow link; MFU collapses. Hence the iron rule: TP degree ≤ GPUs per node (typically ≤ 8). TP is your "shrink one layer below one GPU" tool, and it is a within-node tool only.

2 · Pipeline parallelism — split the layer stack

TP shrinks a layer; it does nothing about a model with more layers than fit. Pipeline parallelism (PP) cuts the L-layer stack into p contiguous stages, one group of layers per device (or per node). A batch flows stage 0 → stage 1 → … → stage p−1 for the forward, then back through for the backward. The only communication is the activation tensor handed from one stage to the next — a single point-to-point send of shape (B, T, d), no collective. That makes PP comm tiny and, crucially, tolerant of slow links: it is the axis you run across nodes.

The problem PP introduces is idleness. If you push one batch through naively, stage 1 sits idle until stage 0 finishes, stage 2 waits on stage 1, and so on — at any instant only one of p stages is busy. The fix is to chop the batch into m micro-batches and stream them, so stages fill up like an assembly line. But the line still has a startup and drain: the first micro-batch must traverse all p stages before the pipeline is full, and the last must drain out. That dead time is the bubble:

The intuition: there are p−1 stage-times of fill and p−1 of drain spread over a schedule of m + p − 1 stage-times of length — so the wasted fraction shrinks as you raise m. With p = 4 stages and m = 4 micro-batches the bubble is 3/7 ≈ 43% — brutal. Raise to m = 32 and it falls to 3/35 ≈ 8.6%; at m = 64, 3/67 ≈ 4.5%. The rule of thumb is m ≳ 4p to keep the bubble in the ~15–20% range; for under ~10% you need m ≳ 8–10p. The cost of more micro-batches is more in-flight activations to store, so the bubble trades against activation memory — another instance of the course's spend-the-cheap-resource move.

1F1B (one-forward-one-backward) is the schedule everyone ships. GPipe runs all m forwards and only then all m backwards, so every micro-batch's activations are alive at the peak — memory scales with m. 1F1B instead, after a short warmup, alternates a forward and a backward on each stage, so a micro-batch's activations are released as soon as its backward runs. Same bubble fraction, dramatically lower peak activation memory — which is what lets you crank m high enough to make the bubble negligible. Interleaved 1F1B goes further, giving each device several non-contiguous stage chunks to shave the bubble further at the cost of more point-to-point sends.

3 · Sequence and context parallelism — split along T

TP splits the hidden dimension d and the heads; PP splits the layer index. Neither helps when the sequence T is the thing that overflows — a 128k-token context produces activations ∝ B·T·d·L and an attention term ∝ B·h·T² that no amount of head-splitting or stage-splitting shrinks, because each token's activation still lives somewhere.

Sequence parallelism (SP) is a cheap companion to TP: the parts of a block that TP leaves replicated and elementwise — the LayerNorms, dropout, the residual adds — are split along T across the same TP ranks, turning TP's all-reduces into all-gather + reduce-scatter pairs of the same total volume while cutting the activation memory of those regions by t. It is nearly free and ships with Megatron TP by default.

Context parallelism with ring attention tackles the attention itself. Shard the sequence into c chunks, one per GPU; each GPU holds its chunk's Q, K, V. To compute attention, the K,V chunks are passed around a ring: GPU i attends its local Q to the K,V it currently holds, computes a partial using the same online-softmax running (max, sum) recurrence from FlashAttention (lesson 7), then forwards that K,V block to its neighbor and receives the next — after c steps every Q has seen every K,V, and the T×T score matrix was never materialized anywhere. Communication (one K,V block per step, point-to-point) overlaps with the local attention compute, so for long T the ring is nearly free. This is how frontier models train at 128k–1M token contexts without any single device holding the whole sequence.

4 · The 3D mesh — composing DP × TP × PP

These axes are orthogonal, so you multiply them. Arrange all your GPUs into a logical 3D mesh: every GPU has a coordinate (dp, tp, pp), and the total must factor exactly: G = DP · TP · PP. Each axis uses the collective it can afford on the link it's placed on. The placement rules fall straight out of the comm costs derived above:

A worked 64-GPU layout. Take 8 nodes of 8 H100s each, 64 GPUs total, training a model too big for any one node. A standard factorization is DP 4 × TP 8 × PP 2 = 64:

Read it as a hierarchy: the 8 GPUs in a node are one TP group (one layer's matmuls split 8 ways on NVLink); two nodes chain into a 2-stage pipeline (16 GPUs holding the full depth, activations handed over IB); and four of those 16-GPU model-groups form a 4-way data-parallel cluster, each chewing a quarter of the global batch and all-reducing gradients at the step boundary. Every GPU holds 1/(TP·PP) of the parameters and 1/DP of the batch. The rules of thumb: fill TP first up to GPUs/node, then PP across nodes to fit depth, then put all remaining GPUs into DP for throughput — and prefer more DP over more PP when the model already fits, because DP comm overlaps and PP adds a bubble.

Interaction with FSDP. The DP axis does not have to be plain replication — it is exactly where ZeRO/FSDP from lesson 8 plugs in. Run ZeRO-1/2/3 along the DP dimension: the optimizer state, gradients, and (for ZeRO-3) parameters are sharded across the DP ranks of the mesh, gathered just-in-time, while TP and PP split each replica's compute. This composition — FSDP as the outer dimension, TP×PP as the inner model split — is the standard frontier recipe: TP shrinks the layer below one GPU, PP fits the depth across nodes, and FSDP-DP both scales throughput and shards the leftover 16-bytes/param state across the widest, cheapest axis. (The same cluster topology shows up again in post-training; see reinforcement_learning · the RL/post-training topology lessons for how rollout, reward, and trainer workers are laid out across a similar mesh.) The fused-optimizer step that runs on each rank is detailed in gpu_kernels · 30 (fused optimizers & the training step).

5 · Drive the mesh planner

Set total GPUs, GPUs per node, model size N, and layer count L, then pick a DP × TP × PP factorization. The planner checks it multiplies to the total, computes the per-GPU memory (state sharded by TP·PP·DP under FSDP), the PP bubble at a sensible micro-batch count, and flags the comm-cost regime. The configuration that breaks: push TP past GPUs/node and the plan turns red — your tensor-parallel all-reduces would cross the node boundary onto InfiniBand and the run would crawl. Prove three things: (a) TP=16 with 8 GPUs/node is invalid; (b) at a fixed GPU count the 16N state is sharded across all three axes, so per-GPU memory depends only on the product DP·TP·PP — raising PP (and dropping DP to compensate) inflates the bubble while memory holds flat; (c) the corollary — shifting that split back from PP into DP lowers the bubble at no memory cost, which is why you prefer DP over PP once the model fits.

The planner makes the whole hierarchy legible at a glance: the only configurations that stay green are the ones where TP fits inside a node, PP carries the cross-node depth, and DP soaks up the rest — exactly the placement the comm costs of §1–§4 forced.

Failure modes & checklist

Checkpoint

Where this points next

You can now train a model of any size: TP shrinks a layer below one GPU, PP fits the depth across nodes, context parallelism handles arbitrary T, and FSDP-DP shards the leftover state across the widest axis. There is no longer a dense model too big to train. But notice what you are paying for that reach — the budget identity C ≈ 6ND from lessons 0 and 05 says every token's forward and backward run through every parameter. Capability scales with total parameters; cost scales with the same total parameters, because each one costs FLOPs on every token. The two are welded together, and the budget caps both at once. The obvious question — the one that breaks this lesson's victory — is whether you can decouple them: grow the parameter count without growing the FLOPs each token pays. You can, by making each token use only a sparse slice of the parameters. That is 10 · Mixture of Experts — parameters without the FLOPs.

Interview prompts

• Why does a column-parallel then row-parallel linear pair need only one all-reduce? (§1 — column-split A gives each GPU a local activation slice; GELU/attention-per-head are elementwise/independent so no comm; row-split B makes the block output the sum of partials, summed by a single all-reduce. Two matmuls, one collective.)
• Why must tensor parallelism stay within a node? (§1 — it does ~4 synchronous all-reduces of ≈B·T·d per layer on the critical path; NVLink is ~900 GB/s vs IB ~25–50 GB/s, a 20–40× gap, so crossing the node boundary stalls every block and MFU collapses. Rule: TP ≤ GPUs/node.)
• State the pipeline bubble fraction and how to shrink it. (§2 — (p−1)/(m+p−1) with p stages, m micro-batches; raise m (rule of thumb m ≳ 4p): p=4,m=4 → 43%, m=32 → ~9%. 1F1B keeps the bubble the same but cuts peak activation memory so you can afford large m.)
• How do tensor and pipeline parallelism differ in communication, and where do you place each? (§1–4 — TP: heavy synchronous all-reduce, NVLink, intra-node; PP: tiny point-to-point activation hand-off, latency-tolerant, across nodes. DP/FSDP outermost as the cheap overlappable axis.)
• What does ring attention buy you, and how does it relate to FlashAttention? (§3 — context parallelism shards T across c GPUs and rotates K,V around a ring, using FlashAttention's online-softmax (max,sum) recurrence to accumulate partials so the T×T scores never materialize; comm overlaps compute → arbitrary-length contexts.)
• Lay out 64 GPUs (8/node) for a model too big for one node. (§4 — DP 4 × TP 8 × PP 2: TP 8 inside each node on NVLink, PP 2 chains two nodes into a pipeline over IB, 4 such model-groups form DP with FSDP sharding the state. TP first to GPUs/node, PP across nodes for depth, DP for the rest.)
• How does FSDP fit into the 3D mesh? (§4 — ZeRO/FSDP runs along the DP axis: optimizer state, grads, and (ZeRO-3) params are sharded across DP ranks and gathered just-in-time, while TP×PP split each replica's compute — FSDP outer, model-split inner.)
• You can train any dense size now — so what still forces a change? (Where this points next — capability ~ total params but cost ~ 6ND runs through every param, so the budget caps both together; decoupling them requires each token touch only a sparse slice of weights → MoE, lesson 10.)