Lab 01 — Distributed Training & Parallelism Engine

Phase: 10 — Distributed Training & Model Parallelism Difficulty: ⭐⭐⭐☆☆ (the algorithms are small; the systems intuition is ⭐⭐⭐⭐⭐) Time: 4–5 hours

One GPU cannot hold a frontier model — recall Phase 00's ~16 bytes/param, so a 70B model needs ~1.1 TB of training memory, an order of magnitude past any single accelerator. The answer is to split the work across many GPUs, and there are exactly four ways to cut it: replicate the model and split the data (DP), split a layer's matrices (TP), split the stack of layers (PP), and shard the optimizer/grad/param state itself (ZeRO/FSDP). This lab builds a runnable model of all four — the ring all-reduce that glues data parallelism together, the column/row matmul sharding of tensor parallelism, the pipeline bubble, and the ZeRO memory ladder — so you can reason about a cluster the way Phase 00 taught you to reason about a single GPU.

What you build

• ring_all_reduce / all_reduce_cost — the collective at the heart of distributed training: reduce-scatter then all-gather around a ring, 2(N-1) steps, ~2·tensor·(N-1)/N bytes/rank — bandwidth-optimal and independent of N as N grows. The result equals a naive element-wise sum (the soul test); the schedule is the lesson.
• data_parallel_gradient_average — the DP update: all-reduce the gradients, divide by N, and prove it equals the serial mean (the entire correctness guarantee of data parallelism).
• tensor_parallel_matmul — Megatron-style sharding of a matmul: column-parallel (concatenate partial outputs, no comm) and row-parallel (all-reduce partial sums). Both must reconstruct the unsharded W · x — the TP soul test.
• pipeline_bubble_fraction / pipeline_schedule_1f1b — the (P-1)/(M+P-1) bubble that pipeline parallelism pays for fill+drain, and the 1F1B step count showing more micro-batches shrink it toward zero.
• zero_memory_per_gpu — the ZeRO/FSDP memory ladder: stages 0→1→2→3 shard the optimizer, then gradients, then parameters across DP ranks, dropping per-GPU bytes from 16N to 16N/N.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why the ring all-reduce is 2(N-1) steps and why bytes/rank approaches 2·tensor (not N·tensor) as ranks grow — the reason it scales.
• You can explain why data_parallel_gradient_average must equal the serial mean, and what breaks if a rank's gradient is dropped or double-counted.
• You can state, for tensor_parallel_matmul, why column-parallel needs no communication and row-parallel needs an all-reduce — and why that all-reduce is what pins TP to one NVLink node.
• You can compute pipeline_bubble_fraction(8, 4) in your head and explain why the fix is more micro-batches, not more stages.
• You can read the ZeRO ladder and say why stage 3 (FSDP) total → 16N/N and what it costs (an all-gather of params every layer).

How this maps to the real stack

Limits of the miniature (be honest in the interview): we model bytes and steps, not wall-clock — real all-reduce overlaps with backward compute, ring latency has a per-step cost the model ignores, and NCCL picks tree vs ring by message size. Tensor parallelism here is a single matmul, not a whole transformer block (the real f/g conjugate all-reduces wrap attention and MLP). The pipeline model counts unit steps, not the real per-stage compute imbalance or activation memory. ZeRO accounting ignores activations (where checkpointing lives) and the all-gather compute/comm cost stage 3 pays to reconstruct each layer's params on the fly.

Extensions (build these on real hardware)

• Add tree_all_reduce_cost(n_ranks, tensor_bytes, latency_per_step) and find the message size where ring beats tree (latency- vs bandwidth-bound) — the real NCCL crossover.
• Wrap tensor_parallel_matmul into a two-layer MLP (column then row) and show the single all-reduce per block, matching Megatron's f/g placement.
• Add interleaved_bubble_fraction(P, M, v) for the v-virtual-stage interleaved schedule and show it divides the bubble by v.
• Add activation_memory_bytes(...) + checkpointing and compose it with zero_memory_per_gpu to size a real 70B training run on 8×80 GB.
• Run the real thing: torchrun --nproc_per_node=2 a tiny DDP/FSDP script, log all_reduce time with torch.profiler, and check it against your all_reduce_cost.

Interview / resume

• Talking points: "Why is the ring all-reduce bandwidth-optimal and 2(N-1) steps?" "Column vs row tensor parallelism — which needs communication and why does TP stay inside one node?" "What is the pipeline bubble and how do you shrink it?" "Walk me up the ZeRO stages and the memory each one buys." "Map a 3D-parallel job onto an 8-node cluster."
• Resume bullet: Built a from-scratch distributed-training engine modeling the ring all-reduce (reduce-scatter + all-gather), data/tensor/pipeline parallelism, and the ZeRO/FSDP memory ladder — reproducing NCCL's 2(N-1)-step cost, Megatron's column/row matmul sharding, the GPipe bubble (P-1)/(M+P-1), and the per-GPU memory drop from 16N to 16N/N.