Parallelism 2

Transcript: cleaned auto-captions with timestamps

The systems arc of CS336 has one recurring question: compute lives far from the data it needs, so how do you arrange the work to keep the arithmetic units fed? Lecture 5 answered it inside a single GPU (registers, shared memory, HBM); lecture 6 built FlashAttention out of that answer; lecture 7 introduced the vocabulary for the answer across GPUs. Lecture 8 is the one where the vocabulary becomes code you can run. Percy's argument is deliberately unglamorous: the multi-GPU strategies that sound intimidating in a systems paper are, at their core, a handful of collective calls dropped into an ordinary training loop, and you understand them best by writing the twenty-line version before you touch Megatron-LM or FSDP.

Outline, with timestamps

• 00:05 — From inside one GPU to across many: same bottleneck, one level up the hierarchy.
• 02:42 — The hierarchy laid out: L1/shared memory → HBM → NVLink → NVSwitch, fast-and-small to slow-and-large.
• 04:21 — Collective operations: world size, rank, and the seven primitives, plus the mnemonic that keeps them straight.
• 08:16 — The wire underneath: PCIe and Ethernet versus NVLink and NVSwitch; H100 numbers; nvidia-smi topo -m on the class cluster.
• 13:10 — NCCL detects the topology and launches the kernels; torch.distributed wraps it, with gloo as the CPU fallback.
• 16:17 — Four processes running one function: spawn, init_process_group, barrier.
• 18:24 — The collectives on four-element tensors: all-reduce, reduce-scatter, all-gather, and the identity that links them.
• 23:47 — Benchmarking all-reduce: warm up, synchronize, barrier, then count bytes. ~273 GB/s.
• 28:13 — Benchmarking reduce-scatter: no factor of two, ~72 GB/s printed, and an honest "I'm not sure why."
• 32:39 — Data parallelism: shard the batch, all-reduce the gradients. The whole diff is one line.
• 40:23 — Tensor parallelism: shard the width, all-gather the activations after every layer.
• 46:18 — Pipeline parallelism: shard the depth, chop the batch into micro-batches, send/recv down the chain.
• 51:39 — What the toy leaves out: async sends, overlap, real bookkeeping — then the Jax/TPU detour and Levanter.
• 61:08 — Summary and questions: batch norm, FSDP versus declarative sharding, DeepSeek, Cerebras, where the graph lives.

One bottleneck, one level up

The opening move is to refuse to treat multi-GPU training as a new subject. Inside a single GPU you have arithmetic units that are fast and data that is far away, and every trick — fusion, tiling, recomputation — exists to keep the units busy while the data catches up. Across GPUs the same sentence is true with different constants. Percy draws the whole memory system as one hierarchy, ordered small-and-fast to large-and-slow, and points out that the reasoning transfers wholesale even though the mechanics do not.

Two things fall out of that table. First, HBM is still roughly four times faster than the best GPU-to-GPU link on the same node, so a byte that crosses to a neighbour costs about what four HBM reads cost — the same "is this trip worth it?" calculus as last week, just re-priced. Second, the gap between the modern path and the commodity path is not a tuning difference. A gaming rig moves data GPU → PCIe → CPU kernel → buffer → back out, and across boxes it hits Ethernet, a protocol designed for a completely different workload. NVLink and NVSwitch exist to delete both hops. 12:37 shows what that looks like in practice: nvidia-smi topo -m on the class cluster prints NV18 in every off-diagonal cell of an 8×8 matrix — every GPU is 18 links away from every other GPU, no PCIe in the path.

Seven primitives and one identity

Collective operations are not new and are not deep-learning-specific — they are 1980s parallel-programming vocabulary, and the reason to use them instead of hand-rolled point-to-point messaging is that a library gets to pick the routing. Two words carry the whole namespace. World size is how many devices are participating; rank is which one you are, an integer from 0 to world size − 1 (nothing to do with matrix rank, a collision Percy flags immediately at 04:52).

The operations decompose along three axes, which is the mnemonic worth memorising: reduce means an associative and commutative combine (sum, min, max, mean); broadcast/scatter are the inverses of gather; and the prefix all- means every rank is a destination, not just one. Run those three switches and you get the table: broadcast sends one rank's tensor everywhere; scatter deals different slices to different ranks; gather is scatter run backwards into one rank; reduce is gather that combines instead of concatenating; all-gather is gather with everyone as destination; reduce-scatter is reduce whose result is split across ranks.

The identity that matters is all-reduce = reduce-scatter + all-gather. It is not a curiosity. It is why all-reduce costs what it costs: a ring implementation moves each rank's payload once around to combine, then once around to redistribute, so the bytes on the wire scale as 2(N−1)/N of the tensor per rank rather than N× it. It is also the seam that ZeRO and FSDP cut along — if you were going to all-reduce gradients and then everyone was going to apply the same update anyway, you can stop halfway, let each rank own a shard of the optimizer state, and all-gather parameters only when you need them. Percy demonstrates the identity concretely at 21:04 by feeding the reduce-scatter output straight into an all-gather and landing on the same tensor all-reduce produced.

NCCL, torch.distributed, and four processes pretending to be one

NCCL is NVIDIA's collective communication library: it discovers the hardware topology at init, picks paths between GPUs, and turns "all-reduce this tensor" into CUDA kernels that push packets. torch.distributed is the Python skin over it, and its useful property is that the backend is swappable — nccl on GPUs, gloo on CPUs. That is not a footnote for a student: it means you can debug the assignment's distributed code on a laptop with no GPU, verify the logic, and only then care about wall-clock. Portability is what you buy by programming against collectives rather than sockets.

The execution model takes a minute to internalise, and the lecture's harness makes it explicit: a helper spawns world_size processes all running the same function, distinguished only by their rank argument. Every process calls init_process_group pointed at a shared coordination address (in the lecture, localhost:15623) — that rendezvous is for discovery only; the tensors themselves never go through it. From there the code reads like single-process code with rank-dependent indexing. dist.barrier() is used in the lecture mostly to keep print statements from interleaving, but the deeper point lands during Q&A: collectives are themselves synchronisation points, which is both how ranks stay in lockstep and how a bug becomes a hang.

"All-reduce is a synchronization point. This will stop everyone and perform an all-reduce operation. So you have to be careful, because if one of your ranks misses the all-reduce, everything will just hang."— Percy Liang, 38:05

That is the single most common failure mode in distributed training code, and it is worth carrying as a rule: any control flow that can differ across ranks — an early return, a data-dependent if, a ragged last batch — is a potential deadlock, because one rank will reach the collective and the others will not.

Measuring it: 273 GB/s, 72 GB/s, and the accounting

The benchmarking segment is the most transferable part of the lecture, and it is mostly about discipline. The recipe: allocate the tensor, run the operation once as a warm-up, call torch.cuda.synchronize() so the CUDA kernels have actually finished, call dist.barrier() so every rank is at the same line, then start the timer, do the operation, synchronize and barrier again, stop the timer. Skip the synchronize and you time kernel launches instead of kernel work; skip the barrier and you time your own rank's head start.

With 100·1024² fp32 elements — 400 MiB per rank — and a world size of 4, all-reduce takes about 2.1 ms. Turning that into a bandwidth number requires deciding what counts as a byte on the wire, and the lecture's convention is: each rank ships its payload to the other N−1 ranks, times two because all-reduce both sends inputs and receives outputs, divided by world_size × elapsed to charge the whole machine for the whole time. That lands at roughly 273 GB/s against a 900 GB/s NVLink rating — a perfectly ordinary result, and exactly the reason Percy insists on measuring rather than quoting the spec sheet.

Reduce-scatter is where it gets instructive. It uses the same recipe but drops the factor of two, because reduce-scatter only moves data one way — there is no redistribution phase. The printed figure is about 72 GB/s, and Percy declines to explain the gap:

"I guess you could say that all-reduce usually has more traffic, and all-reduce is probably more optimized. […] There's a lot going on in NCCL, so it's a little hard to speculate about performance accurately."— Percy Liang, 30:21

The honest reading is that the two numbers are not comparable, because they are not counting the same thing. The all-reduce case moves a 400 MiB tensor; the reduce-scatter case is handed a world_size × num_elements input, which is 1.6 GiB per rank, and takes about 3.9 ms — four times the input in under twice the time. Whether the printed GB/s figure charges for 400 MiB or for 1.6 GiB changes the answer by 4×, and none of that difference is on the wire. Treat collective bandwidth numbers the way you would treat a FLOPs count: reproducible only alongside the formula that produced it, which is why the lecture links NVIDIA's own nccl-tests performance doc rather than asserting a canonical convention.

Three cuts through the same MLP

The second half runs on a deliberately boring model: a four-layer MLP, batch 128, hidden dimension 1024, GELU between layers, loss = mean squared activation. The justification is real — in a Transformer the feed-forward blocks, not attention, dominate the FLOPs at typical shapes, so an MLP is a fair stand-in for the compute you are trying to parallelise. Each strategy is then a different axis to cut the same rectangle.

Data parallelism cuts the batch. With world size 4 each rank takes rows [rank·32 : (rank+1)·32], holds a full copy of every parameter, runs forward and backward on its own slice, and then — this is the entire distributed part — all-reduces each param.grad with ReduceOp.AVG before optimizer.step(). Losses differ across ranks because the data differs; gradients become identical because of the all-reduce; therefore parameters stay identical forever, provided you seeded them identically to begin with. That last clause is a real precondition, not a formality. Note also what is not shared: each rank keeps its own AdamW state, redundantly, which is pure waste — and is precisely the waste ZeRO and FSDP were invented to reclaim. Percy addresses the obvious "why not communicate the optimizer state instead?" at 39:52: updating optimizer state locally is far cheaper than shipping it, so redundant compute beats communication here.

Tensor parallelism cuts the width. Now the data stays whole and the layers get sliced: with num_dim 1024 and world size 4, each rank holds a 1024×256 slab of every layer — one quarter of the parameters, which is the point, since the motivating problem is a model that will not fit on one device. Each rank computes a 128×256 activation slab, then all_gather collects all four slabs and torch.cat along dim 1 reassembles the full 128×1024 activation, which is the input to the next layer. Notice the cost profile: a collective per layer, on activations, in both directions. This is exactly why lecture 7 said tensor parallelism wants a fast interconnect and generally stays inside a node — it is the strategy whose communication scales with the batch and the depth, not with the parameter count.

Pipeline parallelism cuts the depth. Four layers, two ranks, two layers each; every rank sees all the data. The naive version is one rank working while the rest idle — the bubble from lecture 7 — so the batch of 128 is split into 4 micro-batches of 32, and now the primitives change character: dist.recv(tensor=x, src=rank-1) and dist.send(tensor=x, dst=rank+1), point-to-point rather than collective. Rank 0 seeds the chain with real data; every later rank pre-allocates a buffer and fills it from its predecessor. The Q&A around 53:55 covers the sharp edges honestly: sends from one rank to one destination stay ordered, a send nobody receives just blocks, and matching is by source rank rather than by any tensor identity.

What the toy leaves out

Percy is unusually direct about the gap between the lecture code and a real implementation, and the list is worth memorising because it is the list of things frameworks actually sell you. Communication and computation never overlap here: the sends are synchronous when they should be isend calls returning handles you wait on at the end, so the wire sits idle during matmuls and vice versa. The backward pass is skipped for tensor and pipeline parallelism — with pipelining, getting it right means interleaving forward and backward micro-batches (the 1F1B schedules from the GPipe and PipeDream lines of work), which is where most of the real complexity lives. And generalising past a hand-rolled MLP to arbitrary architectures is a bookkeeping problem: FSDP has to introspect parameters and figure out module boundaries, whereas here Percy simply decided how to cut the model.

The Jax detour at 58:20 is the useful contrast. In the Jax/TPU world you declare a model, declare a sharding — "shard this dimension across these devices" — and the compiler emits the collectives. Percy shows FSDP in about ten lines using Levanter, Stanford CRFM's Jax library, and is candid that this is the better abstraction for getting work done; the course stays in PyTorch specifically so you can see the primitives. At the far opposite end he points at DeepSeek, working with a comparatively poor interconnect and hand-tuning below the NCCL level to claw back performance. The spectrum — declare it and let the compiler decide, versus hand-write the packets — is a real engineering choice, and where you land on it is mostly a function of how far your hardware is from the happy path.

The closing synthesis is the sentence to carry out of the systems arc: every parallelism strategy is a choice of which dimension to cut (batch, width, depth, sequence length, expert), and every one of them is a trade among three ways to obtain a value you need — recompute it, store it in local memory, or store it on someone else's GPU and communicate. Activation checkpointing is the first trade; multi-GPU parallelism is the third. The hardware keeps improving, and it never rescues you:

"While they can grow significantly, there are physical limitations and we will always build larger models that are at the limit of the hardware's capabilities."— Percy Liang, 62:54

What you build with this

This lecture is the direct preparation for the distributed half of Assignment 2: Systems (handout PDF), which opened with lecture 5 and whose first half — Triton/FlashAttention kernels and profiling — came from lectures 5 and 6. Percy says it plainly at 39:20: DDP is what you implement in the second assignment, in the context of a Transformer rather than a four-layer MLP, and the lecture code is the stripped-down version so you can see the shape clearly.

Concretely, the pieces this lecture hands you: (1) the spawn / init_process_group / barrier scaffolding, including the gloo backend so you can iterate without a GPU; (2) the benchmarking discipline — warm-up, cuda.synchronize, barrier, explicit byte accounting — which the assignment asks you to apply to all-reduce across backends, device types and payload sizes; (3) naive DDP as a gradient all-reduce, then the optimisation the assignment actually grades, which is batching those per-parameter all-reduces into buckets so you pay fewer, larger collectives and can overlap them with the backward pass; and (4) optimizer state sharding — the ZeRO stage-1 idea that falls straight out of the reduce-scatter/all-gather identity in this lecture. Tensor and pipeline parallelism are not assignment deliverables; they are here so the taxonomy from lecture 7 stops being abstract, and the skipped backward passes are explicitly left as exercises.

Supporting materials, verified

• Executable lecture: lecture_08 trace — Percy Liang (2025) · the lecture itself; step through the same code with live variable inspection.
• lecture_08.py source — stanford-cs336 (2025) · the actual script, including the byte-accounting formulas and all three parallelism implementations. Read this before the video if you prefer code first.
• torch.distributed documentation — PyTorch · the API surface the lecture uses: all_reduce, reduce_scatter_tensor, all_gather_into_tensor, send/recv, backend selection.
• NCCL User Guide — NVIDIA · what the library does underneath, including topology detection and the ring/tree algorithms behind the identity in this lecture.
• NVLink and NVSwitch — NVIDIA · source of the 18-link / 900 GB/s H100 figure Percy quotes at 10:22.
• Collective operation — Wikipedia · the 1980s parallel-programming lineage; linked directly from the lecture script.
• nccl-tests: how to reason about collective performance — NVIDIA · the canonical treatment of algorithmic versus bus bandwidth. Read this if the 273 versus 72 GB/s gap bothered you.
• all_reduce_bench.py — Stas Bekman, ml-engineering · a production-grade version of the benchmark Percy writes live; the lecture links it as the reference implementation.
• Levanter 1.0 — Stanford CRFM (2023) · the Jax library behind the ten-line FSDP snippet in the detour, and the concrete example of declarative sharding.
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al. (2019) · named at 57:49 as where to go next; the real tensor-parallel implementation this lecture's twenty lines gesture at.
• Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al. (2021) · field map extra · how data, tensor and pipeline parallelism compose into 3D parallelism, with the interleaved 1F1B schedule that removes the bubble Percy leaves in.
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism — Huang et al. (2018) · field map extra · the origin of the micro-batching trick used verbatim in the pipeline section.
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al. (2019) · field map extra · the paper that turns "each rank redundantly holds its own AdamW state" into three stages of sharding; the assignment's optimizer-state sharding is stage 1.
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel — Zhao et al. (2023) · field map extra · the library Percy points at twice and deliberately does not use; explains the bookkeeping that arbitrary-architecture sharding demands.
• PyTorch Distributed: Experiences on Accelerating Data Parallel Training — Li et al. (2020) · field map extra · gradient bucketing and backward/communication overlap — i.e. the exact distance between the lecture's one-line DDP and the real one.
• DeepSeek-V3 Technical Report — DeepSeek-AI (2024) · field map extra · the "opposite end of the spectrum" Percy names at 65:37: training at scale on a constrained interconnect by engineering below the collective-library level.

Exercises

1. Reproduce the bandwidth curve code — take the lecture's all_reduce benchmark and sweep it. (1) Copy the warm-up / cuda.synchronize / barrier / time / synchronize / barrier skeleton. (2) Sweep tensor size from 1 MiB to 1 GiB in powers of four, at world sizes 2, 4 and 8. (3) Plot measured GB/s against payload size. (4) Repeat with the gloo backend on CPU. A good answer shows the small-message region where latency dominates and bandwidth is terrible, the knee, and the plateau — and names roughly where the knee sits on your hardware, because that number is what decides whether gradient bucketing is worth it.
2. Audit the byte accounting — derive, on paper, the bytes-on-the-wire for a ring all-reduce of a tensor of B bytes across N ranks, and show it is 2B(N−1)/N per rank rather than the lecture's 2B(N−1). Then work out which convention the lecture's ~273 GB/s and ~72 GB/s figures each use, and say precisely why the two numbers cannot be compared head-to-head. A good answer states the payload each measurement actually moved (400 MiB versus a 1.6 GiB input) and concludes with the wall-clock comparison, which is convention-free.
3. Finish the backward passes code — the lecture skips backward for tensor and pipeline parallelism and calls both homework. (1) For tensor parallelism, work out what collective the gradient of an all_gather + cat is (hint: it is the operation that is its adjoint) and implement one training step. (2) Verify against a single-process reference by seeding identically and asserting gradients match to 1e-5. (3) For pipeline parallelism, add the reverse send/recv chain and a naive all-forward-then-all-backward schedule. (4) Count the bubble: what fraction of rank-time is idle at 4 micro-batches versus 16? A good answer includes the numeric agreement check, not just code that runs.
4. Bucket the gradients code — the lecture's DDP issues one all-reduce per parameter, which is the slow way. (1) Flatten gradients into fixed-size buckets (try 1 MB, 25 MB, 100 MB). (2) All-reduce per bucket instead of per parameter and unflatten back. (3) Time a step against the naive version on the lecture's 4-layer, 1024-wide MLP scaled up to ~50 layers. (4) Then register backward hooks so a bucket fires as soon as its last gradient is ready, overlapping communication with the rest of the backward pass. A good answer reports the speedup from bucketing alone and from bucketing-plus-overlap separately, and connects the result to the knee you found in exercise 1.