01 — The 5D Parallelism Taxonomy

The mental discipline of this phase: for every parallelism strategy, know (a) what state is sharded, (b) what collective is invoked, (c) how many bytes cross the wire per token, (d) which fabric the collective lands on. If you can produce those four numbers on demand, you understand it.

The five axes

1. Data Parallelism (DP)

Every GPU holds a full copy of the model and processes a different micro-batch. After the backward pass, gradients are all-reduced across replicas.

• State sharded: none (fully replicated).

• Collective: one all-reduce over gradients per optimizer step.

• Bytes per step: 2 × params × bytes_per_grad (ring all-reduce).

• Fabric: anywhere — gradient all-reduce is one-per-step, not per-token, so PCIe or IB is fine.

• When it dies: when model + optimizer states don’t fit on one GPU. Then you sharded state, and DP becomes FSDP/ZeRO.

2. Tensor Parallelism (TP) — Megatron style

Split individual matmuls across GPUs. For a matmul Y = X W:

• Column-parallel: shard W along output dim. Each GPU produces a shard of Y. Followed by an all-gather (or fused into next op).

• Row-parallel: shard W along input dim (requires X already sharded along matching dim). Each GPU produces a partial sum. Followed by an all-reduce.

Megatron pairs them cleverly (see 02_tensor_parallelism.md): the MLP is column-then-row (one all-reduce at the end); attention QKV projection is column, output projection is row (again one all-reduce at the end). So two all-reduces per transformer block in the forward pass, two more in the backward.

• State sharded: weights, activations, KV.

• Collective: all-reduce, per block, per token.

• Bytes per token per block: 2 × hidden_size × bytes in each direction of the ring, twice (attention + MLP), forward + backward.

• Fabric: must be NVLink. TP over PCIe is a joke; TP over IB is a crime. If you don’t have an NVLink domain of size ≥ TP degree, don’t do TP that wide.

• When it dies: past NVLink domain. On H100 SXM (NVSwitch of 8), TP=8 is the max sane value. On H200 NVL72, TP=72 becomes possible — this is Blackwell’s structural advantage.

3. Pipeline Parallelism (PP)

Split layers across GPUs. Layer 0–k on GPU-0, k+1–2k on GPU-1, etc. Micro-batches flow through the pipeline like an assembly line.

• State sharded: weights (by layer partition).

• Collective: point-to-point send/recv at each stage boundary. No global collective per token.

• Bytes per token per boundary: hidden_size × bytes (activation) forward + hidden_size × bytes (grad) backward.

• Fabric: IB or NVLink both fine. This is why PP is cross-node friendly.

• The pain: bubbles — idle time at fill and drain. Scheduling schemes (GPipe, 1F1B, interleaved, zero-bubble) exist to minimize this. See 03_pipeline_parallelism.md.

• When it dies: at small batch (bubbles dominate), and at very deep pipelines (activation memory for stashed micro-batches explodes).

4. FSDP / ZeRO (Sharded Data Parallelism)

A hybrid: like DP in that every rank participates in every layer, but state is sharded across ranks. Before each layer’s forward, the full weights are all-gathered; used; then discarded. Backward all-gathers again for grad computation, then reduce-scatters grads.

• State sharded: weights, gradients, optimizer states (stage 3); grads + optim (stage 2); optim only (stage 1).

• Collective: all-gather per layer (fwd), all-gather + reduce-scatter per layer (bwd).

• Bytes per step per layer: ~2× params_per_layer × bytes across ring.

• Fabric: fast intra-node preferred; can survive IB with prefetch overlap.

• When it dies: communication overwhelms compute overlap. Very small models on very fast networks — the all-gather isn’t hidden by any compute.

See 04_fsdp_zero.md.

5. Expert Parallelism (EP) — MoE only

Each expert (an MLP in the MoE layer) lives on a specific GPU. Tokens are dispatched to their chosen experts via all-to-all, computed, then combined back via all-to-all.

• State sharded: MoE expert weights only. Dense layers are TP/DP as usual.

• Collective: two all-to-alls per MoE layer, per token.

• Bytes per token per MoE layer: 2 × top_k × hidden_size × bytes (dispatch + combine).

• Fabric: wants fastest fabric available. All-to-all is bandwidth-hungry and topology-sensitive.

• Load balance is a first-class problem: if experts see uneven traffic, some GPUs idle. Aux losses (Switch Transformer) or aux-loss-free bias updates (DeepSeek-V3) fix this. See 05_expert_parallelism.md.

6. Context / Sequence Parallelism (CP / SP)

Split the sequence dimension across GPUs. Ring Attention passes K/V blocks around the ring while each GPU computes attention against its local Q block. Enables million-token contexts.

• State sharded: activations along sequence axis, KV cache along sequence axis.

• Collective: ring pass of K, V (send-recv, structured).

• Bytes per iteration: seq_shard × head_dim × heads × bytes × 2 per ring step.

• Fabric: fast. Ring latency compounds.

• When it pays: context ≥ 32k. At 128k+ it’s mandatory for training.

See 06_sequence_parallelism.md.

Composing: 3D and 4D

Real training runs compose these. A canonical 3D setup for a 70B on a 64-GPU cluster:

DP=4, TP=4, PP=4  →  4 × 4 × 4 = 64 GPUs

With MoE and long context you get 4D or 5D:

DP=8, TP=4, PP=4, EP=32, CP=2

Composition order rules:

1. TP innermost (fastest fabric). Group of TP-rank GPUs shares NVLink.

2. EP next (still wants fast fabric — all-to-all is heavy).

3. PP next (cross-node friendly).

4. DP/FSDP outermost (least frequent comms).

5. CP slots into whichever fabric can support the ring — typically same domain as TP.

The torch.distributed.device_mesh and Megatron’s parallel_state APIs express this composition.

Communication cost per token — the actual table

For a transformer block with hidden size h, L layers, TP degree t, PP degree p, DP degree d, in bf16 (2 bytes):

Strategy	Bytes crossing wire per token per block	Frequency
DP	0 (per token); 4h²/L per param per step	1 all-reduce / step
TP	~4h(t-1)/t (two all-reduces per block, ring)	Every block, every token
PP	2h (activation send + grad send)	Every stage boundary
FSDP	~4h²/L · (d-1)/d per param	Every block (all-gather)
EP (MoE)	2 · top_k · h	Every MoE block
CP	2h · seq_shard per ring step	Ring-length iterations

(Coefficients are approximate; the exact factor depends on whether all-reduce is ring or tree and whether reductions are fused. See 07_nccl_networking.md for the ring formulas.)

The one-line takeaway: TP communicates per token; DP communicates per step. That’s a factor of ~batch × seq_len difference. It’s why TP wants 900 GB/s NVLink and DP survives on 400 Gbps IB.

Decision matrix (put this on a sticky note)

Situation	Strategy
Model fits on one GPU, want more throughput	DP (or DDP)
Model + optimizer barely fits	ZeRO-1 or FSDP full-shard
Optimizer state huge (AdamW)	ZeRO-1 minimum
Model doesn’t fit even at bf16	FSDP-3 / ZeRO-3, or TP if you have NVLink
Model > 1 node	TP within node + PP across nodes
MoE model	Add EP; keep dense parts TP
Context > 32k in training	Add CP
Inference, latency-sensitive, batch 1	TP within node
Inference, throughput-focused	Prefer DP replicas of a smaller-parallelism config

Exercises

1. Llama-3-70B, bf16, 8k context, TP=4 on 4×H100. Compute the bytes/second of TP all-reduce traffic at 50 tokens/sec generation and batch 1. Compare to H100 NVLink SXM aggregate bandwidth (900 GB/s).

2. Take a 400B MoE with top_k=2, hidden=8192, 8 GPUs of EP. Compute all-to-all bytes per token per MoE layer.

3. Argue on paper: for a 7B on 8×A100 40GB, is FSDP-3 or TP=8 better for training throughput? (Answer will hinge on activation memory and comm-overlap.)

Reading order for the rest of this folder

02 (TP) → 03 (PP) → 04 (FSDP) → 05 (EP) → 06 (SP) → 07 (NCCL). Then you have the vocabulary for serving (08–10) and the doorway into training (11+).