Lecture 07 · Part 2

Parallelism I

Percy scales the resource-accounting mindset from one GPU to many. The building blocks are the collective operations — all-gather, reduce-scatter, all-reduce — implemented over a hardware hierarchy where NVLink is fast, InfiniBand is slow, and ethernet goes through the CPU. With those primitives he builds bare-bones data, tensor, and pipeline parallelism on an MLP, each cutting the training computation along a different axis.

▶ Lecture video Executable lecture (.py) ↗

The same game, one level up

Last week’s question was how to make one GPU go fast; this week’s is how to use many. The picture extends naturally: inside each box you still have HBM, L2, L1, and streaming multiprocessors, but now there are four boxes, or a thousand, connected by a network. And the situation, Percy stresses, is structurally identical: compute (the ALUs and tensor cores) is far from data. On a single GPU “far away” meant HBM; across GPUs the tensor you need might live on a different device entirely. “It’s very easy to use a ton of GPUs, but it’s hard to use them effectively.”

The memory hierarchy simply grows a few levels. L1 and shared memory are fastest; then HBM — which last week we lamented as slow, but “in this lecture, HBM is going to be considered fast”; then NVLink/NVSwitch within a node; then InfiniBand and ethernet across nodes, slowest of all. Last week’s tricks (fusion, tiling) reduced memory accesses; this week’s (replication, sharding) reduce communication across GPUs.

Why multi-GPU at all? Two reasons. First, capacity: a B200 has 192 GB of HBM, and a 1-trillion-parameter model’s weights, gradients, optimizer state, and activations simply don’t fit. Second, speed: even a model that fits can be trained faster by splitting the work — if the communication you pay doesn’t eat the compute you gained. That trade is the calculation this lecture teaches you to make.

Collective operations: the vocabulary

Four warm-ups, then the workhorses. Broadcast copies a tensor from rank 0 to every rank (used maybe once per run, to distribute an initial checkpoint). Scatter splits a tensor on rank 0 into world-size pieces and sends one to each rank. Gather is its inverse: pieces from all ranks concatenate onto rank 0. Reduce is gather with an associative-commutative operation applied — pieces 0, 1, 2, 3 sum to 6 on rank 0.

The three that drive distributed training add “all”, meaning the destination is all devices. All-gather performs a gather onto every rank: each rank holds a piece, and afterwards every rank holds the concatenation. Reduce-scatter reduces each dimension and scatters the results — if rank $r$ holds the vector $[r, r{+}1, r{+}2, r{+}3]$ for $r = 0, \dots, 3$ , then afterwards rank 0 holds 6 (the sum of the first components), rank 1 holds 10, rank 2 holds 14, rank 3 holds 18. And all-reduce leaves the full reduced vector $[6, 10, 14, 18]$ on every rank.

The most general primitive is all-to-all: each rank sends its $j$ -th piece to rank $j$ , so for balanced splits it acts like a matrix transpose. It matters for mixture-of-experts training, where each rank holds a slice of the data and a subset of experts, and routing is dynamic — you must look at the data to know which expert (hence which rank) each activation goes to. Percy’s mnemonic for the zoo: reduce applies an operation (sum, min, max), scatter is the inverse of gather, and “all” means the destination is all devices.

Hardware: what the network actually looks like

The classic home setup — GPUs hanging off a PCIe bus (v7.0 at 16 lanes: 242 GB/s), machines linked by ethernet at roughly 200 MB/s — is what you’d get wiring two gaming GPUs together. Serious training looks different. Typically 8 GPUs per node connect via NVLink to an NVSwitch: NVLink 5.0 delivers 1.8 TB/s of total bandwidth, about 4x slower than the B200’s 8 TB/s HBM but very fast for inter-device traffic, and the switch means any GPU can talk to any other GPU without you managing the routing. Around 256 nodes form a pod connected by InfiniBand (through PCIe to an InfiniBand NIC, roughly 0.05 TB/s), and pods connect by ethernet, through PCIe and the CPU. The pattern echoes the memory hierarchy: the more devices you span, the slower the link.

NVIDIA keeps pushing the NVLink domain outward: the GB200/GB300 NVL72 racks put 8 GPUs per tray and 9 trays per rack, so 72 GPUs share one NVLink domain instead of the mortal eight. Underneath all of it sits NCCL, the NVIDIA Collective Communications Library: you say “all-reduce”, and NCCL detects the hardware topology, picks paths between GPUs, and launches the actual communication kernels — because everything that runs on a GPU, communication included, is a kernel.

torch.distributed, and what bandwidth you actually get

PyTorch wraps NCCL in torch.distributed, with a gloo backend for CPUs and higher-level machinery (FSDP) that this course deliberately skips — we build from primitives. A program spawns world-size processes, each running the same function with its own rank.

Code walkthrough — The three workhorses in code

Each process first initializes the process group (the master address/port is only for coordination — data goes through NCCL), then issues collectives. barrier() makes all processes wait at a point, since ranks otherwise run fully asynchronously. All-reduce modifies its tensor in place:

data = tensor([0., 1, 2, 3], device=cuda_if_available(rank)) + rank
dist.all_reduce(tensor=data, op=dist.ReduceOp.SUM, async_op=False)
# every rank now holds [6, 10, 14, 18]

Reduce-scatter and all-gather instead take explicit input and output tensors:

input = torch.arange(world_size, dtype=torch.float32, device=...) + rank
output = torch.empty(1, device=...)
dist.reduce_scatter_tensor(output=output, input=input, op=dist.ReduceOp.SUM, async_op=False)

input = output                                  # feed reduce-scatter's result in
output = torch.empty(world_size, device=...)
dist.all_gather_into_tensor(output_tensor=output, input_tensor=input, async_op=False)

Chaining them reproduces the all-reduce result on every rank — “proof via example that all-reduce is equal to reduce-scatter plus all-gather”. Setting async_op=True returns immediately, letting you overlap communication with computation, at the cost of managing the waits yourself.

How fast is this in practice? Benchmark like last week: warm up, then time — but with two forms of asynchrony to flush, so you call torch.cuda.synchronize() (CUDA kernels finish) and dist.barrier() (all processes arrive) before reading the clock.

Worked example — Effective bandwidth of an all-reduce

All-reduce a tensor of $100 \cdot 1024^2$ fp32 elements (about 419 MB) across 4 ranks; it takes roughly 1.6 ms. To turn that into a bandwidth, count bytes that must move: the reduction takes world size $-$ 1 steps, and each step both sends and receives, so

$\text{sent bytes} = \text{size} \cdot 2 \cdot (W - 1), \qquad \text{total duration} = W \cdot \text{duration}$

which gives $419\text{MB} \cdot 2 \cdot 3 \,/\, (4 \cdot 1.6\text{ms}) \approx 400$ GB/s. As $W$ grows, $(W-1)/W \to 1$ , so effective bandwidth $\approx 2 \cdot \text{size} / \text{duration}$ — independent of world size, and independent of whether NCCL chose a ring or tree topology. Reduce-scatter measures out to the same ~400 GB/s with no factor of 2: all-reduce moves twice the data in twice the time, and the twos cancel.

Distributed training: three ways to cut

Part two trains deep MLPs — the compute bottleneck of a Transformer, so representative of the real thing — with sample data of batch size 128 and dimension 1,024, and cuts the computation three different ways.

Data parallelism splits the batch. With world size 4, each rank takes its slice of 32 rows (data[start_index:end_index]), keeps a full copy of all parameters and its own AdamW state, and trains almost normally:

for step in range(num_steps):
    x = data                        # this rank's 32 rows only
    for param in params:
        x = x @ param
        x = F.gelu(x)
    loss = x.square().mean()
    loss.backward()

    # Sync gradients (the ONLY difference between standard training and DDP)
    for param in params:
        dist.all_reduce(tensor=param.grad, op=dist.ReduceOp.AVG, async_op=False)

    optimizer.step()

Tensor parallelism cuts the other way: every rank sees all the data, but each layer’s weight matrix is split down its columns, so rank $r$ holds a num_dim x local_num_dim slice (1024 x 256 here). Each rank computes its slice of the activations, applies the elementwise GELU (safe, since it’s elementwise), then all-gathers so every rank can reassemble the full activation for the next layer:

x = x @ params[layer]              # batch_size x local_num_dim
x = F.gelu(x)
activations = [torch.empty(batch_size, local_num_dim, device=...) for _ in range(world_size)]
dist.all_gather(tensor_list=activations, tensor=x, async_op=False)
x = torch.cat(activations, dim=1)  # back to batch_size x num_dim

Unlike DDP, you now have to muck around with the model — this leans on the fact that a matmul decomposes into smaller matmuls. In the backward pass the all-gather turns into a reduce-scatter of gradients: the two operations are duals across forward and backward.

Pipeline parallelism cuts by depth: each rank takes a contiguous subset of layers (2 of the 4 here, world size 2) and the batch flows through them in stages, using point-to-point dist.recv from rank $-$ 1 and dist.send to rank $+$ 1. The naive version leaves GPUs idle — pipeline bubbles — so the batch is chunked into micro-batches (4 of size 32) that stream through, keeping stages busier. What this bare-bones version skips is overlapping communication with computation, which for pipeline parallelism is not optional garnish but the crucial ingredient — and something assignment 2 explores for the data-parallel case too, where gradients can start flowing as soon as backward produces them.

Percy closes with the recurring pattern: you can recompute (activation checkpointing), store in memory, or now store on another GPU and communicate. DDP does redundant work — every rank updates every parameter — precisely to avoid moving optimizer state. Hardware keeps getting faster, but we’ll always want bigger models, so this hierarchy, and these trade-offs, are permanent.

Back of the chalk tray — what to remember

The multi-GPU game is the single-GPU game scaled up: compute is far from data, and you orchestrate to avoid transfer bottlenecks — replication and sharding instead of fusion and tiling.
Collective operations are the vocabulary: all-gather, reduce-scatter, and all-reduce do the work of training, and all-reduce = reduce-scatter + all-gather.
The interconnect hierarchy — NVLink (1.8 TB/s) within a node, InfiniBand (~0.05 TB/s) within a pod, ethernet beyond — mirrors the memory hierarchy, and RDMA (bypassing the CPU) is what makes the fast tiers fast.
Effective bandwidth of an all-reduce is about 2 x size / duration (~400 GB/s measured), independent of world size and topology.
DDP is standard training plus one all-reduce of gradients per step; tensor parallelism splits layers by columns and all-gathers activations every layer; pipeline parallelism splits by depth and fights bubbles with micro-batches.
Choose parallelism by interconnect: tensor parallel needs NVLink-class bandwidth, pipeline tolerates slow links, data parallel scales until the critical batch size.

Check yourself

Exercises in the lecture’s spirit — collective-operation reasoning and communication napkin math.

Problem 7.1 effective bandwidth

You all-reduce a tensor of $200 \cdot 1024^2$ fp32 elements across 8 ranks and each rank measures 3.5 ms. Compute the effective bandwidth using the lecture’s accounting, and state what it converges to for large world sizes.

Show solution

Size: $200 \cdot 1024^2 \cdot 4$ bytes $= 800$ MiB $\approx 0.839$ GB. Sent bytes: $\text{size} \cdot 2 \cdot (W - 1) = 0.839 \cdot 2 \cdot 7 \approx 11.7$ GB. Total duration: $W \cdot \text{duration} = 8 \cdot 3.5\text{ms} = 28$ ms. Bandwidth $\approx 11.7 / 0.028 \approx 420$ GB/s. For large $W$ , $(W-1)/W \to 1$ , so the expression converges to $2 \cdot \text{size}/\text{duration}$ — independent of world size, which is why it’s a fair figure of merit as you scale the cluster.

Problem 7.2 collective op mechanics

Four ranks; rank $r$ holds the vector $[10r,\ 10r{+}1,\ 10r{+}2,\ 10r{+}3]$ . Write out what every rank holds after (a) a reduce-scatter with sum, then (b) an all-gather of those results. What single operation is the composition?

Show solution

(a) Component $j$ sums to $\sum_{r=0}^{3}(10r + j) = 60 + 4j$ , and reduce-scatter puts component $j$ ‘s reduction on rank $j$ : rank 0 holds 60, rank 1 holds 64, rank 2 holds 68, rank 3 holds 72. (b) All-gather concatenates the shards onto every rank, so all four ranks hold $[60, 64, 68, 72]$ . The composition is exactly an all-reduce with sum — the identity the lecture demonstrates in code, and the seam where ZeRO/FSDP insert themselves.

Problem 7.3 DDP communication cost

A 1B-parameter model trains with DDP in fp32 gradients over an interconnect whose effective all-reduce bandwidth is 400 GB/s. Estimate the gradient-synchronization time per step, and name the technique that hides it.

Show solution

Gradient size: $10^9 \cdot 4 = 4$ GB. Effective bandwidth $\approx 2 \cdot \text{size}/\text{duration}$ , so duration $\approx 2 \cdot 4 / 400 = 0.02$ s — about 20 ms per step, regardless of world size. To hide it, overlap communication with computation: gradients for the last layers are ready first during backward, so start all-reducing them while earlier layers are still backpropagating, rather than waiting for the full backward pass to finish.

Problem 7.4 choosing a parallelism

An $L$ -layer MLP with hidden dimension $D$ trains on batches of $B$ rows across $W$ ranks. Per step, compare the elements communicated by (a) column tensor parallelism’s forward-pass all-gathers and (b) DDP’s gradient all-reduce (count tensor elements moved into the collective, ignoring constant factors). For $B = 128$ , $D = 1024$ , which moves more — and why does tensor parallelism still demand the faster interconnect?

Show solution

(a) Tensor parallel all-gathers the activations at every layer: $B \cdot D$ elements per layer, $L \cdot B \cdot D$ per forward pass. (b) DDP all-reduces every gradient once per step: $L \cdot D^2$ elements. The ratio is $D / B = 1024/128 = 8$ , so DDP moves 8x more elements here. Tensor parallelism still needs NVLink because its communication sits on the critical path inside every layer — each all-gather must complete before the next layer can start, so latency and bandwidth both bite $L$ times per pass — whereas DDP’s single all-reduce happens once per step and can be overlapped with the backward pass.