The setup: two bottlenecks, two speed regimes

We scale to many machines to solve compute (one chip is far below the exaflops of a cluster) and memory (models don't fit in one GPU). The recurring conceptual split is intranode communication (very fast — you can afford expensive patterns) vs internode (slow — you must respect the channel). Everything today is at the level of collective operations, and one identity does a lot of work:

Hardware aside: TPU mesh vs GPU tree

The networking differs by philosophy. A TPU uses a toroidal mesh — each chip wired to neighbors, wrapping around — so it scales indefinitely with a fixed number of neighbors per chip (cheap, beefy neighbor links). A GPU uses a fat tree — fast at the bottom, then pods, then spine switches — an all-to-all philosophy that's more flexible for unpredictable communication. Hence (per Dally/Dean) GPUs suit MoEs (tokens route unpredictably) and TPUs suit dense, predictable partitions (tensor parallel over the mesh).

It's converging, though: newer TPUs moved toward tree/all-to-all (Virgo network) because modern models are MoEs that fling activations around. And the Huawei Ascend 910 shows the opposite extreme — weaker chips but a giant fiber-optic all-to-all rack of 384, at ~4× NVIDIA's power: you can brute-force communication if you'll pay the power (the SRAM/Groq story again). The real takeaway: the new unit of compute is the data center, and we want to control memory and compute losslessly across it.

Naive data parallel — and why memory is the problem

Split the batch B across M machines; each computes gradients on B/M elements; all-reduce to sum. Compute scales perfectly; communication is ~2×params per step. But memory savings: zero — every GPU holds a full copy of everything. And "everything" is a lot:

Rule of thumb: ~16 bytes per parameter (≈5 copies). The optimizer state dominates — which is exactly what ZeRO attacks.

ZeRO: shard the redundant state

Stage 1 — shard optimizer state

Each worker updates one parameter slice. Everyone computes a full gradient → reduce-scatter gradients → each does its slice's update → all-gather updated params. That's reduce-scatter + all-gather = one all-reduce — identical comms to DDP. Free.

Stage 2 — also shard gradients

You can't materialize the whole gradient, but you don't need to: sweep backward and reduce-scatter each layer's gradient to its owner as you go, freeing it immediately. Incremental = same cost. Still free.

Stage 3 (FSDP) — shard params too

Each GPU holds only a slice of params, grads, and state. Gather params on demand while stepping the compute graph: all-gather a layer's weights → forward → free; on backward, all-gather again → compute → reduce-scatter grads → free. That's two all-gathers + one reduce-scatter — one extra all-gather vs DDP.

FSDP is clean, elegant, and you'll implement it in the assignment (wrap a module: all-gather → compute → free, repeat). But it has two limits, which force the rest of the toolbox: it consumes the batch size (you can't have more data-parallel workers than batch elements, and past the critical batch size more batch helps less than more steps), and it doesn't touch activation memory.

Pipeline parallelism — cut the depth

Put different layers on different GPUs; pass activations forward, partial gradients backward. Done naively, utilization is catastrophic — one GPU active at a time, the rest idle in a bubble:

The fix is micro-batches: chop the batch so a stage finishes a piece and passes it on, keeping the pipeline full. Utilization ≈ stages / micro-batches, so the bubble shrinks as 1/microbatches — another use for the precious batch-size resource. Pipeline parallel's virtues: it saves memory, composes with data parallel, and communicates only activations (b·s·h, small) point-to-point — so it tolerates slow links and goes on the slowest interconnects (cross-pod, cross-datacenter, decentralized).

Tensor parallelism — cut the width

Split each matmul into smaller matmuls and recombine (the tiling idea again). For an MLP z = GeLU(xA)·B: split A by columns and B by rows, run in parallel, combine at the end.

forward:   f = copy(x)        →  parallel matmuls  →  g = all-reduce
backward:  g = identity       ←  (partials flow)   ←  f = all-reduce
# column-cut at inputs (MLP up-proj, attention QKV);
# row-cut at outputs (MLP down-proj, attention out);
# layernorm / nonlinearity / MoE router → replicated (don't bother cutting)

The forward/backward duality matters: the f and g operators swap roles (identity ↔ all-reduce) between passes. Tensor parallel is very communication-hungry — an activation-sized all-reduce at every matmul — so it lives only on the fastest interconnects (within a node, ≤8 GPUs on NVLink); cross-node it falls off a cliff. (TPUs' big mesh is the exception — they tensor-parallel over far more chips.)

Activation memory — the part everyone forgets

Memory isn't just parameters. Profiling shows huge activation humps, and at scale activations dwarf parameter memory. Storing everything costs roughly:

activation_memory ≈ 34·s·b·h  +  5·(a·s/h)·s·b·h
#                  └ the fundamental s·b·h term ┘   └ quadratic attention+dropout ┘
#  (the second term is droppable via FlashAttention recomputation)

The s·b·h dependence is fundamental — you store something per sequence position, per batch element, per hidden dim. How each strategy reduces it:

Tensor parallel divides the matmul parts (24 of the 34) and attention by t, but the lightweight ops it replicates — layernorm, dropout, the residual inputs you must keep — leave a stubborn 10·sbh penalty. Sequence parallel fixes that: split those leftover ops along the sequence axis (all-gather / reduce-scatter on demand — FSDP-like, with the same forward/backward duality). Add recomputation and you reach 34·sbh/t — the practical lower bound to remember when checking whether a model fits.

Expert parallelism (and why it's preferred for MoE)

Shard the experts (the FFNs) across devices — analogous to tensor parallel (high-bandwidth, reduces activation), but you route sparse tokens instead of slicing dense matmuls. For MoE, prefer EP over TP (per Megatron's guidelines): tensor parallel makes matmuls too small (hurting GPU utilization), whereas routing sparse token activations is cheaper, lets you send tokens exactly where needed, and skips overhead. But EP is genuinely hard — heavy all-to-all dispatch that's latency-sensitive (compute waits for tokens). DeepSeek's DeepEP and NVIDIA's hybrid-EP libraries optimize this at the lowest level — DeepSeek even found undocumented PTX instructions to squeeze out networking performance.

Context parallel (ring attention) splits a long sequence's activations across accelerators, passing them ring-like along the mesh — used for long-context extension and serving.

No strategy dominates — so stack them

You can do the math per layer (compute vs communication time) and plot utilization: as long as compute time exceeds comms time, comms hides underneath. With a big batch, FSDP alone is compute-bound and great; as batch shrinks it becomes comms-bound, so you add tensor parallel to push the efficient region out, then more strategies — that stacking is 3D/4D parallelism.

The classic Narayanan/Megatron sweeps bear this out: DP maxed, TP rising to 8 then capped, pipeline parallel filling the rest, DP shrinking only at extreme scale — and utilization stays flat and high even at enormous GPU counts. (A counterintuitive lever: do more activation recomputation — it saves memory, which buys a bigger batch, which buys better utilization.)

Real training runs

Common thread: everyone maximizes data parallel, keeps tensor parallel ≤ 8, and — thanks to DeepSeek-V3-era infrastructure — expert parallel can now be large. NVIDIA's Megatron-Bridge repo publishes recommended configs per model.

Golden rules & takeaways

• ZeRO shards the redundant state. Stages 1–2 (optimizer state, gradients) are communication-free; stage 3 (FSDP, params too) adds one all-gather that overlap hides — near single-GPU utilization at a fraction of the memory.
• FSDP's limits force model parallelism: it consumes the batch size and ignores activation memory.
• Model parallelism communicates activations. Pipeline (depth, point-to-point, slow-link-friendly, fight bubbles with micro-batches), tensor (width, all-reduce-heavy, fast-link-only), expert (shard experts, all-to-all, preferred for MoE).
• Activations dominate memory at scale; tensor + sequence parallel + recomputation drive it toward 34·sbh/t.
• No strategy dominates — stack 4D. Cut to fit with TP/EP (intranode) + PP/ZeRO-3, then data-parallel the rest; keep TP ≤ 8, pipeline across slow links, EP for MoE, CP for long context.
• The data center is the unit of compute — and clever parallelism keeps utilization flat even across data centers.

Next: scaling laws.

References

Primary source is the lecture itself. Key works named:

1. Rajbhandari et al. (2019) — ZeRO
2. Zhao et al. (2023) — PyTorch FSDP
3. Shoeybi et al. (2019) — Megatron-LM (tensor parallelism)
4. Narayanan et al. (2021) — Efficient Large-Scale Training on GPU Clusters (3D parallelism)
5. Korthikanti et al. (2022) — Reducing Activation Recomputation (sequence parallel)
6. Qi et al. (2024) — Zero Bubble Pipeline Parallelism
7. Liu et al. (2023) — Ring Attention (context parallelism)
8. NVIDIA Megatron-Core / Megatron-Bridge — parallelism guides & configs