Phase 10 — Distributed Training & Model Parallelism

The phase where one GPU stops being enough and you learn to think in clusters. Phase 00 taught you that training costs ~16 bytes/param, so a 70B model needs ~1.1 TB of memory and a frontier run burns 6ND FLOPs that no single accelerator can finish this decade. The senior question is no longer "how big is the model" but "how do I cut it across 8 / 64 / 4096 GPUs, and what does the wiring between them cost?" That cut — data, tensor, pipeline, or sharded — is the difference between a model that trains and one that OOMs on line one.

Why this phase exists

Every frontier model is trained on a cluster, and the cluster is not a faster GPU — it is a network. The moment you split a model, the bottleneck moves from FLOPs to communication: the gradients that must be summed across data-parallel replicas, the activations that must hop between pipeline stages, the partial sums that must be all-reduced inside a tensor-parallel layer. A Senior AI Engineer treats the interconnect (NVLink inside a node, InfiniBand between nodes) as a first-class budget, exactly the way Phase 00 treated HBM bandwidth. Almost every distributed decision is downstream of a few facts you can compute on paper:

1. The collective. The ring all-reduce is 2(N-1) steps and moves ~2·tensor·(N-1)/N bytes per rank — bandwidth-optimal and ~constant in N. This one algorithm is why data parallelism scales at all.
2. Data parallelism. Replicate the model, split the batch, all-reduce the gradients — the averaged gradient is provably the serial mean. Cheap and the default, until the model itself won't fit on one GPU.
3. Tensor parallelism. Shard a layer's matrices (Megatron column/row). Reconstructs the exact forward, but pays an all-reduce inside every block — so it lives inside one NVLink node.
4. Pipeline parallelism. Split the stack of layers across GPUs; pay the bubble (P-1)/(M+P-1) for fill and drain, and shrink it with more micro-batches (GPipe → 1F1B).
5. ZeRO / FSDP. Don't replicate the optimizer/grad/param state — shard it across the DP ranks. Stages 1→2→3 drop per-GPU memory from 16N to 16N/N, which is how a 70B model trains on commodity 80 GB cards at all.

Get fluent here and "train a 175B model" stops being magic and becomes DP × TP × PP × ZeRO mapped onto a topology — a design you can draw on a whiteboard.

Concept map

                         ┌──────────────────────────────────────┐
                         │  One GPU can't hold it: ~16 B/param   │
                         │  70B train ≈ 1.1 TB  →  split it       │
                         └──────────────────────────────────────┘
                          split DATA │ split a LAYER │ split LAYERS │ shard STATE
                ┌───────────────┘        │             │              └────────────┐
                ▼                         ▼             ▼                           ▼
          DATA PARALLEL            TENSOR PARALLEL  PIPELINE PARALLEL          ZeRO / FSDP
        replicate model           shard W (col/row) split the stack         shard opt→grad→param
        all-reduce grads          all-reduce in-block  bubble (P-1)/(M+P-1)  16N → 16N/N per GPU
                │                         │             │                           │
                └────────── glued by ─────┴──────┬──────┴───────────────────────────┘
                                                 ▼
                                   COLLECTIVE COMMUNICATION
                          all-reduce = reduce-scatter + all-gather
                            ring: 2(N-1) steps, ~2·tensor bytes/rank
                                                 │
                                                 ▼
                                       3D PARALLELISM (DP × TP × PP)
                         TP inside a node (NVLink) · PP across nodes · DP outermost
                              overlap comm with compute · map to the topology

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py.

Integrated scenario ideas

• Size a 70B training run: with ~16 bytes/param it needs ~1.1 TB — show that ZeRO-3 on 16×80 GB drops it to ~70 GB/GPU (just barely), and what activation checkpointing buys on top.
• Diagnose a stalled scale-out: a DP job that scaled to 8 GPUs goes slower at 64. Compute the all-reduce bytes/rank and show it's network-bound; propose gradient bucketing + overlap, or a hierarchical all-reduce, not "more GPUs."
• Place 3D parallelism on a cluster: given 8 nodes × 8 GPUs, choose TP=8 (inside NVLink), PP=8 (across nodes), DP=1 — and explain why TP must not cross the node boundary.
• GPipe vs 1F1B: for P=8, show the bubble at M=8 (44%) vs M=64 (8%), and why 1F1B's memory footprint beats vanilla GPipe at the same bubble.
• DDP vs FSDP for a 13B model: compute per-GPU memory under DDP (won't fit) vs FSDP FULL_SHARD (fits), and name the cost FSDP pays (an all-gather of params each layer).

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive the ring all-reduce step count 2(N-1) and the ~2·tensor·(N-1)/N bytes/rank from the reduce-scatter + all-gather decomposition.
• You can explain why data-parallel gradient averaging equals one big batch, and what an all-reduce actually computes.
• You can state, for tensor parallelism, which sharding (column/row) needs communication and why TP stays inside one node.
• You can compute a pipeline bubble fraction and say how to shrink it.
• You can walk up the ZeRO stages and quote the per-GPU memory each one saves.

Key takeaways

• The interconnect is the budget. Once you split a model, communication — not FLOPs — is the bottleneck. The ring all-reduce is bandwidth-optimal precisely because each rank moves ~2·tensor bytes regardless of cluster size; understand that and you understand why DP scales.
• Four cuts, one cluster. Data parallelism for throughput, tensor parallelism for a layer too big for one GPU (inside NVLink), pipeline parallelism for a stack too deep (across nodes), and ZeRO/FSDP to stop replicating the 16N state. Real runs compose all four — 3D parallelism.
• ZeRO/FSDP is the memory unlock. Replicating the optimizer state (12N of the 16N) on every GPU is the waste; sharding it across N ranks turns 16N into 16N/N. That single idea is why frontier models train on commodity 80 GB cards.
• Pipelines have a bubble; micro-batches pay it down. (P-1)/(M+P-1) → 0 as M → ∞, and 1F1B gets the same bubble at a fraction of GPipe's activation memory. More stages is not free.

Next: Phase 11 — RAG & Vector Search.