Parallelism 1

Transcript: cleaned auto-captions with timestamps

The systems arc of CS336 has been climbing a ladder of abstraction. Lecture 5 was about what a single GPU actually does with a matmul; lecture 6 was about writing kernels that keep it busy. Lecture 7 is where the unit of compute stops being a chip and becomes a building. That change is not cosmetic: once your model lives on more than one device, the thing you are optimising is no longer arithmetic throughput but a budget split across memory, wires, and batch size — and the wires are wildly non-uniform. Everything in the rest of the lecture falls out of one fact about the hardware, so Tatsu spends the first ten minutes on the network before he shows you a single algorithm.

Outline, with timestamps

• 00:04 — Framing: from one GPU to the datacenter: compute and memory both run out; multi-node is not optional.
• 03:22 — The hardware hierarchy: NVLink/NVSwitch inside a box, InfiniBand between boxes, another cliff past ~256 GPUs.
• 05:03 — Collectives refresher, and the identity that pays for everything: all-reduce = reduce-scatter + all-gather.
• 08:28 — GPU all-to-all vs TPU toroidal mesh; the datacenter becomes the unit of compute.
• 12:29 — Three axes to cut along: data, model, activation parallelism.
• 14:09 — Naive data parallel, and the 16-bytes-per-parameter memory bill that sinks it.
• 18:47 — ZeRO stage 1: shard the optimizer state, pay nothing extra in bandwidth.
• 27:37 — ZeRO stages 2 and 3 (FSDP): shard gradients, then parameters, and hide the comms behind compute.
• 43:04 — Batch size is a finite resource, and it is where data parallel stops.
• 45:18 — Pipeline parallel: cutting by depth, the bubble, 1F1B, and zero-bubble scheduling.
• 55:31 — Tensor parallel: cutting by width, two all-reduces per block, and why you stop at 8.
• 63:21 — Activation memory: the term nothing else shards, and sequence parallel as the fix.
• 71:06 — Ring attention and expert parallel in passing; the comparison table.
• 76:41 — 3D parallelism rules of thumb, the Megatron-LM scaling table, and what real 2024 models do.

The hierarchy is the argument

An H100 node is eight GPUs bolted to a switch. Inside that box, any GPU can talk to any other over NVLink through an NVSwitch at enormous bandwidth. Step outside the box and you are on InfiniBand, roughly an order of magnitude slower per lane — Tatsu puts it at about 8× slower, and the exact multiplier matters less than the fact that there is a cliff. Step outside the rack, past roughly 256 GPUs, and you hit a second cliff, where leaf and spine switches replace the all-to-all fabric (09:04).

Hold onto that shape, because every recommendation later in the lecture is a consequence of it. Bandwidth-hungry algorithms go inside the box. Bandwidth-frugal algorithms go across boxes. There is no cleverness in the rule; the cleverness is in figuring out which algorithm is which.

Google's TPUs make a different bet. Instead of an all-to-all fabric with a cliff, they wire chips to their nearest neighbours in a toroidal mesh: you can only talk to your neighbours, but the mesh extends indefinitely without a discontinuity. Tatsu's point is that this is not a downgrade, because collectives like all-reduce and reduce-scatter are exactly the operations a ring or torus implements optimally. If the only thing you ever do is collective communication — and in distributed training it very nearly is — a torus is a perfectly good topology, and it means TPU shops feel much less pressure to reach for pipeline parallelism (50:18).

Before the algorithms, one identity, which Tatsu flags as the thing you must carry with you: an all-reduce is exactly a reduce-scatter followed by an all-gather, and in the bandwidth-limited regime the decomposition costs the same as the fused operation — about 2× the size of the reduced object either way. That equivalence looks like trivia. It is the entire reason ZeRO stage 1 is free.

"Instead of the GPU, the new unit is the data center."— Tatsunori Hashimoto, 10:15

Data parallel, and the memory bill nobody budgets for

Naive data parallelism is the obvious move: split the batch B across M machines, let each compute gradients on its slice, all-reduce the gradients, step. Compute scales linearly. Communication is 2× the parameter count per batch, which is fine when batches are large enough to amortise it. And memory scales not at all — every GPU still holds a full copy of everything.

"Everything" is worse than it sounds. Tatsu walks the accounting for mixed-precision AdamW training and lands on roughly 16 bytes per parameter, in about five copies of the weights: 2 bytes of BF16 parameters, 2 bytes of BF16 gradients, 4 bytes of FP32 master weights, and 4 bytes each for Adam's first and second moment estimates. The model parameters themselves — the thing you actually care about — are 2 of those 16 bytes. Optimizer state dominates.

Which suggests an obvious question: does every rank really need every optimizer state? The ZeRO family answers no, in three stages, and the payoff is dramatic. For a 7.5B model over 64 GPUs, the per-device memory goes 120 GB → 31.4 GB → 16.6 GB → 1.9 GB as you shard optimizer state, then gradients, then parameters (18:47).

Stage 1 shards only the optimizer state. Each rank still computes a full gradient on its own data slice, but only updates the parameter shard it owns. The mechanics: reduce-scatter the gradients so each rank ends up with the fully-summed gradient for its slice, take the Adam step locally, then all-gather the updated parameters back out. Count the traffic — one reduce-scatter plus one all-gather — and you get exactly the 2×-params of a naive all-reduce. This is the identity from ten minutes earlier cashing in. Stage 1 gives you a large memory win for zero extra bandwidth, which is why Tatsu's advice is that you may as well always turn it on.

Stage 2 also shards gradients. The complication is that you can never materialise a full gradient vector, so the reduction has to be interleaved with the backward pass: as each layer's gradient is produced, reduce it immediately to its owning rank and free it. More synchronisation points, but the same 2×-params total.

Stage 3 — this is what PyTorch calls FSDP — shards parameters too. Now no rank holds a full layer, so parameters must be all-gathered on demand as you walk the graph forward and backward, and freed immediately after use. The traffic goes to 3× params: two all-gathers plus one reduce-scatter. A 1.5× bandwidth increase for full linear memory scaling is a good trade, but the genuinely surprising part is the latency story. Because the all-gather for layer n+1 can be issued while layer n is still computing, the communication hides behind the compute almost entirely; Tatsu's walkthrough of the FSDP timeline (34:10) is the clearest thing in the lecture, and it is essentially prefetching applied to weights.

The practical numbers, on a single 8×A100-80GB node, assuming pure BF16 training with FP32 master weights only (12 bytes per parameter rather than 16):

One more virtue worth naming: ZeRO is architecture-agnostic. Nothing in it knows what a transformer is. That is why FSDP ships as a block wrapper you can drop over an arbitrary nn.Module, and why it is the default for anyone who does not want to build training infrastructure.

Batch size is a resource, and it runs out

Data parallel has a hard ceiling that has nothing to do with memory or bandwidth: you cannot have more data-parallel ranks than you have examples in the batch. You cannot put a third of an example on a GPU.

"You can't parallelize greater than your batch size, because you can have at most one example on each machine."— Tatsunori Hashimoto, 43:10

And you cannot simply raise the batch size to buy more parallelism, because past a point extra examples stop helping optimisation. This is the critical-batch-size result: below it, your gradient is noisy and averaging more examples genuinely reduces variance; above it, you are limited by the number of steps you take, not by gradient quality, and the marginal example buys almost nothing.

The reframe Tatsu wants you to make is that global batch size is a budget. Data parallel spends it. Pipeline parallel also spends it, because micro-batches are what fill the pipeline bubble. Tensor parallel is the one strategy that spends none of it — which turns out to be a major reason it is used so aggressively despite being the most bandwidth-hungry option on the list.

Cutting the model: depth versus width

Model parallelism differs from ZeRO-3 in a way worth stating precisely, because a student asks exactly this at 40:48. Both shard parameters. But ZeRO-3 ships parameters to wherever the compute is; model parallelism leaves parameters where they are and ships activations instead. When activations are smaller than parameters, that is a much better deal — and it also shards activation memory, which ZeRO-3 does not.

Pipeline parallel cuts along depth: layers 1–8 on GPU 0, 9–16 on GPU 1, and so on, with activations handed forward and gradients handed back. Done naively this is the worst parallelism imaginable — with n GPUs each is busy 1/n of the time and you have bought n devices for the throughput of one. The fix is micro-batching: split the batch into n_micro pieces and keep them in flight, so the idle "bubble" shrinks to a fraction (n_stages − 1) / n_micro of useful compute. That formula is the whole story — pipelines only work with large batches, and NVIDIA's ablations show per-GPU utilisation collapsing at batch size 8 and holding up nicely at 128.

Why tolerate it? Two reasons. It shards activations as well as parameters, and its communication is point-to-point and proportional to activations only — roughly b·s·h per micro-batch — which makes it the natural choice for the slow links between nodes and racks. The frontier of pipeline scheduling is zero-bubble scheduling (DeepSeek calls their variant DualPipe), which exploits the observation that the backward pass has two separable halves: propagating gradients with respect to activations (serially dependent) and computing gradients with respect to weights (dependent on nothing downstream, so schedulable anywhere). Slot the weight-gradient work into the holes in a 1F1B schedule and the bubble largely disappears.

The catch is that implementing this means reaching inside autodiff and running your own work queue, and Tatsu is unusually blunt about the human cost:

"There's two people in the group that understand how the pipeline parallel in our infra works. One person left. And so there's a single load-bearing person in our training infra."— Tatsunori Hashimoto, relaying a frontier-lab anecdote, 54:56

Tensor parallel cuts along width instead. Since almost all the parameters and FLOPs live in matmuls, split the matrices: give A₁ and B₁ to one GPU and A₂, B₂ to another, compute partial products, and sum. In the Megatron MLP layout, the input x is broadcast to both ranks, each computes GeLU(xAᵢ)Bᵢ, and the results are all-reduced. The two synchronisation points f and g are duals: in the forward pass f is identity and g is an all-reduce; in the backward pass the roles swap.

Conceptually simpler to implement than pipelining, no bubble, no batch-size cost — but it synchronises twice per transformer block and moves roughly 8·b·s·h·(n−1)/n per layer via all-reduce. That is a lot of traffic, which is why it lives and dies on interconnect speed. The empirical cliff Tatsu shows from Hugging Face's parallelism guide is stark: 10–12% throughput loss up to 8-way TP, 42% at 16-way, 65% at 32-way (59:18). Eight is not a magic number in the algorithm — it is the number of GPUs behind one NVSwitch.

Activations: the term that refuses to shard

All of the above is about parameter-shaped memory. Activations are a separate, dynamic budget, and the memory-over-time plot Tatsu shows (63:56) makes the point well: peak usage lands partway through the backward pass, where activations have not all been freed and gradients are already accumulating.

Per transformer layer, storing everything, activation memory is

s·b·h · (34 + 5·a·s/h)

where s is sequence length, b batch, h the residual width and a the head count. The left term is MLP and pointwise work; the right term is the quadratic attention scratch, and — since the h cancels — it grows as s². FlashAttention-style recomputation kills that second term outright, which is a satisfying callback to lecture 6.

Apply tensor parallel everywhere it applies and you get s·b·h·(10 + 24/t). That stubborn 10·s·b·h is the residue: LayerNorm, dropout, and the inputs to attention and the MLP. None of them are matmuls, so tensor parallel does not touch them, and they grow with model size forever.

Sequence parallelism is the one-line fix, and it is almost embarrassingly simple once stated: LayerNorm and dropout are pointwise across positions, so shard them along the sequence axis instead. The seams need collectives — all-gather where you re-enter tensor-parallel territory, reduce-scatter where you leave it, reversed on the backward pass — but the payoff is that the last un-sharded term finally divides by t. Combine sequence parallel with recomputation and per-layer activation memory reduces to 34·s·b·h / t, which is the figure you will see quoted in transformer-arithmetic cheat sheets.

Two more axes get a passing mention. Context (ring) parallel shards activations along a long sequence, passing keys and values machine-to-machine in a ring while each rank owns a slice of queries — which works precisely because you already know attention can be computed tile-by-tile from FlashAttention. Expert parallel scatters MoE experts across devices; conceptually it is tensor parallel, except routing is sparse and unbalanced, so the communication pattern is far less predictable.

Putting it together, and what labs actually do

The summary table is worth internalising as four rows, each dominated by a different cost:

And the rule of thumb, which Tatsu notes has not changed between his 2024 and 2025 versions of this slide:

The 2021 Megatron-LM scaling study is the empirical backbone here. Across models from 1.7B to 1T parameters they hold 40–52% of theoretical peak FLOPs, and the configuration table reads exactly like the rule above: tensor parallel climbs to 8 and stops; pipeline parallel then grows to make the model fit; data-parallel size falls from 32 to 6 as pipeline stages eat the batch. Tensor-parallel 8 comes out optimal at both large and small batch sizes, and an 8×8 TP/PP grid wins across 64 devices. Activation recomputation also pays for itself: it costs FLOPs but buys the larger batches that hide the pipeline bubble — the same trade FlashAttention makes one lecture earlier.

The 2024 model reports line up. OLMo/Dolma at 7B is plain FSDP, because it fits inside a node. DeepSeek's first model is ZeRO-1 plus tensor, sequence and pipeline parallel — the vanilla combination. DeepSeek-V3 is the interesting outlier: 16-way pipeline, 64-way expert parallel across 8 nodes, ZeRO-1 for the data dimension, and no conventional tensor parallel. Yi is ZeRO-1 + TP + PP; Yi-Lightning replaces tensor with expert parallelism because it is an MoE. Llama 3 405B uses TP 8, plus context parallel for the long-context stage only, plus pipeline and data parallel — and the report's stated ordering of strategies by bandwidth appetite is the same ordering Tatsu derived from first principles.

The Llama 3 footnote worth carrying is not about parallelism at all. At that scale, 148 interruptions came from faulty GPUs — about 30% of all interruptions, with another 32 from unplanned maintenance. Tatsu adds that practitioners consider silent data corruption scarier than outright failure, because a GPU that quietly returns garbage can poison a run without announcing itself. Fault tolerance is not a footnote to the algorithms; at 16k GPUs it is a peer of them.

What you build with this

This lecture is the direct spec for the back half of Assignment 2: Systems (handout PDF, leaderboard). Parts 1 and 2 — benchmarking, profiling, and the FlashAttention-2 Triton kernel — come from lectures 5 and 6. Parts 3 and 4 are this lecture, almost problem for problem.

Part 3, distributed data parallel. You benchmark the collectives themselves (distributed_communication_single_node), implement naive DDP with a per-parameter all-reduce (naive_ddp), then attack the two inefficiencies Tatsu names: too many small communication calls, fixed by flattening (minimal_ddp_flat_benchmarking), and waiting for the whole backward pass before communicating, fixed by firing all-reduces from post-accumulate-grad hooks as each gradient lands (ddp_overlap_individual_parameters) and then bucketing them (ddp_overlap_bucketed). That overlap is the FSDP timeline from 34:10, implemented by hand. Tests: tests/test_ddp.py, tests/test_fsdp.py.

Part 4, optimizer state sharding. You write a wrapper around an arbitrary torch.optim.Optimizer that assigns each rank ~1/world_size of the parameters, steps only its own shard, and broadcasts updated parameters afterwards — ZeRO stage 1, exactly as diagrammed at 21:38. Then you profile peak memory before and after the optimizer step and check the accounting against the 16-bytes-per-parameter breakdown. Test: tests/test_sharded_optimizer.py.

The written problem communication_accounting is where the lecture's arithmetic gets exercised directly: an XXL config (d_model 16384, d_ff 53248, 126 blocks), FP32 master weights and optimizer state, BF16 activations and gradient communication, and questions about how many H100s of memory that is, what FSDP degree gets you under one TPU v5p's 95 GB, and at what per-device batch size the forward pass becomes compute-bound. Do it with the Scaling Book's formulas open beside you.

Supporting materials, verified

• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari, Rajbhandari, Ruwase, He (2019/2020) · The source of stages 1–3 and of the 120 → 31.4 → 16.6 → 1.9 GB figure Tatsu quotes for a 7.5B model on 64 devices.
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel — Zhao et al. (2023) · Cited directly on the FSDP slide; the paper behind the communication/computation overlap timeline.
• Getting Started with Fully Sharded Data Parallel (PyTorch tutorial) — PyTorch docs · The source of the "baby version" FSDP diagram; the fastest path from this lecture to running code.
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al. (2019) · Where the tensor-parallel MLP and attention split, and the f/g duality, come from.
• Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al. (2021) · The scaling table (1.7B → 1T, 40–52% of peak), the TP-8 optimality result, and the batch-size-versus-bubble ablations.
• Reducing Activation Recomputation in Large Transformer Models — Korthikanti et al. (2022) · The activation-memory formulas, the residual 10·s·b·h term, and sequence parallelism. Cited by name on the slides.
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism — Huang et al. (2018) · The micro-batching idea that turns layer-wise partitioning into a usable pipeline.
• Zero Bubble Pipeline Parallelism — Qi, Wan, Huang, Lin (2023) · The backward-pass split (activation gradients vs weight gradients) that lets you fill the bubble; up to 23% over 1F1B at comparable memory.
• Ring Attention with Blockwise Transformers for Near-Infinite Context — Liu, Zaharia, Abbeel (2023) · Context parallelism: keys and values circulate in a ring while each device owns a slice of queries.
• An Empirical Model of Large-Batch Training — McCandlish, Kaplan, Amodei et al., OpenAI (2018) · The critical-batch-size result behind "batch size is a finite resource".
• How to Scale Your Model, Part 5: How to Parallelize a Transformer for Training — Austin, Douglas, Frostig, Levskaya et al., Google (2025) · The "TPU book". Source of the batch-size-per-device phase diagram showing when FSDP alone suffices and when you must mix in model parallelism; also the reference the assignment's communication-accounting problem expects you to use.
• GPT-NeoX-20B: An Open-Source Autoregressive Language Model — Black et al. (2022) · Tatsu attributes the eight-GPU node diagram to this paper, with a hedge; the topology it illustrates is the one that governs every rule in the lecture.
• OLMo: Accelerating the Science of Language Models — Groeneveld et al., AI2 (2024) · The 7B run Tatsu cites as plain FSDP. (The slide labels it "Dolma", which is the dataset paper; the training setup is in the OLMo paper.)
• DeepSeek LLM: Scaling Open-Source Language Models with Longtermism — DeepSeek-AI (2024) · ZeRO-1 plus tensor, sequence and pipeline parallel — the textbook combination.
• DeepSeek-V3 Technical Report — DeepSeek-AI (2024) · 16-way pipeline, 64-way expert parallel across 8 nodes, ZeRO-1; also the paper that introduces DualPipe, the zero-bubble variant Tatsu names.
• Yi: Open Foundation Models by 01.AI — 01.AI (2024) · ZeRO-1 + tensor + pipeline parallel.
• Yi-Lightning Technical Report — 01.AI (2024) · An MoE that swaps tensor parallelism for expert parallelism, with pipeline and context parallel alongside.
• The Llama 3 Herd of Models — Meta AI (2024) · TP 8 + CP + PP + DP, ordered by bandwidth appetite; also the failure statistics (148 faulty-GPU interruptions, ~30% of the total).
• Gemma 2: Improving Open Language Models at a Practical Size — Google DeepMind (2024) · The TPU case study Tatsu closes on: ZeRO-3 plus model and data parallelism, with model parallelism stretched further than a GPU fabric would allow.
• The Ultra-Scale Playbook: Training LLMs on GPU Clusters — Tazi et al., Hugging Face (2025) · Field map extra. The most detailed free treatment of exactly this material, with interactive memory and throughput calculators; the assignment handout points at its pipeline-parallel appendix.
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training — Narayanan et al. (2018) · Field map extra. Where the 1F1B schedule Tatsu uses as the starting point for zero-bubble comes from.
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding — Lepikhin et al., Google (2020) · Field map extra. The origin of expert parallelism and its all-to-all routing, which lecture 4 covers from the modelling side and this lecture only gestures at.
• PyTorch Tensor Parallelism (DTensor) documentation — PyTorch · Field map extra. The current production API for the Megatron-style split described at 55:31, including the sequence-parallel variants.

Exercises

1. Redo the memory table from scratch — Without looking at the slide, write the bytes-per-parameter for baseline, ZeRO-1, ZeRO-2 and ZeRO-3 under two precision regimes: (a) BF16 params + BF16 grads + FP32 master weights + FP32 Adam moments, and (b) pure BF16 with FP32 master weights only. For each, compute the largest model that fits on 8×A100-80GB and on 8×H100-80GB. A good answer states the per-GPU formula symbolically — something of the form p + g/N + (m + v + w)/N — before plugging in numbers, and notes explicitly that these figures exclude activations, which is why the answer is optimistic in practice.
2. Bubble and bandwidth arithmetic — You have 512 H100s (64 nodes of 8), a model that needs 6 TB of parameter-shaped memory, and a critical batch size of 4M tokens at sequence length 8192. Choose TP, PP and DP degrees. Show: the bubble fraction (n_stages − 1)/n_micro for your pipeline depth; the per-layer tensor-parallel all-reduce volume 8·b·s·h·(t−1)/t against NVLink bandwidth; and the ZeRO-1 gradient traffic against InfiniBand. A good answer arrives at TP=8, states which of the three budgets binds first, and says what it would change if the model were half the size.
3. Naive DDP → overlapped DDP code — Against assignment2-systems on 2 GPUs: (1) implement naive_ddp with one all-reduce per parameter after loss.backward(); (2) benchmark it and record the fraction of step time spent in communication; (3) flatten all gradients into one buffer and all-reduce once; (4) move to register_post_accumulate_grad_hook so each gradient is communicated the moment it exists; (5) bucket the hooks so you issue a handful of large all-reduces instead of hundreds of small ones, and sweep bucket size. Verify with uv run pytest tests/test_ddp.py at each stage. A good answer includes a table of step time versus bucket size and explains the U-shape — small buckets pay per-call latency, huge buckets lose overlap because nothing can be sent until late in the backward pass.
4. ZeRO stage 1 by hand, and prove the accounting code — Implement the sharded-optimizer wrapper from part 4 of the assignment: partition parameters across ranks in add_param_group, step only your shard, broadcast updated parameters. Then instrument it: report peak memory after model init, immediately before the optimizer step, and immediately after, with and without sharding, using torch.cuda.max_memory_allocated and a memory snapshot. A good answer shows measured optimizer-state memory dropping by roughly 1/world_size while parameter and gradient memory stay flat, matches that against the predicted bytes-per-parameter, and reports the wall-clock overhead of the extra broadcast — which should be small, because per Tatsu the traffic is identical to a plain all-reduce.