Lecture 8: Parallelism 2

What this lecture covers

This lecture follows directly from the previous one's conceptual overview of parallelism, and grounds those ideas in code. It walks through how collective communication operations (broadcast, scatter, gather, reduce, all-gather, reduce-scatter) are implemented in NVIDIA's NCCL library and exposed through PyTorch's torch.distributed, benchmarks their bandwidth on real hardware, and then builds bare-bones implementations of data, tensor and pipeline parallelism on a simple deep MLP.

After watching, you should be able to write basic distributed PyTorch code using collective operations, explain what happens on the wire when you call all_reduce or reduce_scatter, and see concretely how data, tensor and pipeline parallelism cut a model and its data along different dimensions.

Key ideas

• Hardware hierarchy: GPUs contain streaming multiprocessors with a small fast L1 cache, larger high-bandwidth memory (HBM), and links to other GPUs (NVLink within a node, NVSwitch/InfiniBand across nodes); newer hardware bypasses the CPU and Ethernet to connect GPUs directly, which is much faster than the traditional PCIe-and-Ethernet path.
• NCCL and torch.distributed: NCCL translates collective operations into low-level GPU-to-GPU data transfers and topology-aware routing; torch.distributed wraps this in a portable Python interface that also supports a CPU backend (Gloo) for debugging without a GPU.
• World size and rank: world size is the number of participating devices/processes, and rank is the index (0 to world size minus one) identifying each one.
• Benchmarking collectives: measuring bandwidth requires accounting for how much data each rank actually sends; an all-reduce transfers roughly twice as much as a reduce-scatter or an all-gather alone, matching the earlier lecture's claim that all-reduce equals reduce-scatter plus all-gather.
• Data parallelism in code: each rank gets a distinct slice of the batch, runs an ordinary forward and backward pass, then calls all_reduce on each parameter's gradient before the optimizer step, synchronizing parameters across ranks while losses can differ.
• Tensor parallelism in code: each rank holds a slice of every layer's hidden dimension, and after each layer an all-gather reassembles the full activation before the next layer, which is why tensor parallelism needs high-bandwidth interconnects.
• Pipeline parallelism in code: each rank holds a contiguous block of layers; ranks pass activations forward with point-to-point send/receive, and splitting a batch into micro-batches reduces (without fully removing) the idle-time bubble.
• Recomputation as a general tradeoff: choosing to recompute a value rather than store or transfer it recurs at every scale, from single-GPU activation checkpointing to multi-GPU memory and communication tradeoffs.

Walkthrough

Hardware hierarchy and why it matters (0:05)

The lecture reviews the memory/communication hierarchy from L1 cache up to cross-node links, and explains that modern GPU clusters connect GPUs directly via NVLink and NVSwitch rather than routing through the CPU and Ethernet, because that traditional path adds overhead unsuited to deep learning workloads.

Collective operations reviewed (4:08)

It restates the collective primitives (broadcast, scatter, gather, reduce, all-gather, reduce-scatter) with a mnemonic: 'reduce' means an associative combining operation like sum, 'scatter' is the inverse of gather, and 'all' means every device receives the result.

NCCL and torch.distributed in practice (13:14)

The lecture shows a live walkthrough of a multiprocessing utility that spins up several ranks running identical code, initializing a process group, then calling all_reduce, reduce_scatter and all_gather on toy tensors to show exactly what values land where after each call.

Benchmarking collective operations (23:22)

Using a 100-million-element tensor across four ranks, it times all_reduce and computes achieved bandwidth, noting the factor-of-two data movement in all-reduce versus reduce-scatter, and comparing the measured bandwidth against the hardware's theoretical peak to emphasize that real numbers depend on tensor size, device count and other factors.

Coding data parallelism (32:36)

A four-layer MLP is trained with each rank taking a different slice of the batch; the only change from single-device SGD is inserting an all_reduce on each parameter's gradient before the optimizer step, which keeps parameters identical across ranks even though local losses differ.

Coding tensor parallelism (40:43)

Each rank holds a fraction of every layer's hidden dimension. After computing a local slice of the activations, ranks perform an all-gather to reconstruct the full activation before the next layer, illustrating why tensor parallelism is communication-heavy and needs fast interconnects.

Coding pipeline parallelism (45:48)

Layers are split across ranks, and each rank receives activations via point-to-point communication, applies its layers, and sends the result to the next rank. Splitting the batch into micro-batches is shown as a way to shrink the pipeline bubble, though this implementation is synchronous and omits the backward pass and communication/computation overlap.

Beyond PyTorch: JAX and hardware tradeoffs (58:00)

The lecture closes with a brief look at JAX's declarative sharding on TPUs, contrasted with the more manual, low-level tuning some teams do directly against NCCL on GPU clusters, and a discussion of physical limits and specialized inference hardware such as Groq and Cerebras.

Before you watch

• Watch the previous lecture (Parallelism 1) first, since this one assumes familiarity with data, tensor and pipeline parallelism concepts and the all-reduce/reduce-scatter/all-gather identity.
• Basic familiarity with PyTorch and writing a training loop is assumed.

Check your understanding

1. Why does measuring the bandwidth of an all-reduce require multiplying by a factor of two, while reduce-scatter does not?
2. In the data-parallel code example, why do the ranks compute different losses but end up with identical parameters?
3. What communication step does tensor parallelism add after each layer, and why does that make it bandwidth-hungry?
4. How does splitting a batch into micro-batches help pipeline parallelism, and what does this implementation still leave unoptimized?
5. What is the general tradeoff between recomputation, storing in memory, and communicating across GPUs?