Lecture 7: Parallelism 1

What this lecture covers

This lecture asks how you train language models that no longer fit on a single GPU, or that would take too long to train on one. It moves the unit of computation from a single GPU up to a whole data center, and works through the parallelization strategies used to split both compute and memory across many machines: data parallelism, model parallelism (pipeline and tensor), and activation/sequence parallelism.

It follows directly from the prior lecture's single-GPU optimization and builds toward understanding how real frontier labs train large models. After watching, you should be able to reason about the memory and communication cost of each parallelism strategy, know when each is appropriate given hardware constraints, and understand how they are combined into the 3D/4D parallelism used in practice.

Key ideas

• Collective communication primitives: all-reduce, broadcast, reduce, all-gather and reduce-scatter are the building blocks parallelization algorithms are implemented from; reasoning about performance reduces to counting these operations.
• All-reduce equals reduce-scatter plus all-gather: this identity has the same bandwidth cost but lets you interleave computation between the two steps, which is the basis of ZeRO.
• ZeRO / FSDP: shards optimizer state (stage 1), gradients (stage 2), and parameters (stage 3, equivalent to FSDP) across GPUs in data parallelism, cutting memory use roughly in proportion to the number of GPUs at increasing communication cost (free, 2x, then 3x parameter communication).
• Batch size as a resource: batch size is a finite quantity that can be spent on data parallelism, on hiding pipeline bubbles, or on gradient accumulation; it cannot be split below one example per device.
• Pipeline parallelism: splits the model by layer across devices and passes activations forward and gradients backward; naive versions leave most GPUs idle (a 'bubble'), which micro-batching and schedules like zero-bubble/dual-pipe reduce but do not eliminate.
• Tensor parallelism: splits individual matrix multiplies across devices, requiring an all-reduce per layer; it needs very fast interconnects, so it is typically limited to the roughly eight GPUs within one machine.
• Sequence parallelism: splits point-wise operations (layer norm, dropout) along the sequence dimension to shrink the activation memory that tensor parallelism alone cannot reduce.
• Combining strategies: a common rule of thumb is tensor parallel within a machine, then FSDP or pipeline parallel across machines until the model fits in memory, then data parallel to use remaining GPUs.

Walkthrough

Why multi-machine parallelism is needed (0:05)

The lecture opens by framing the problem as both a compute and a memory constraint: single GPUs cannot keep up with model growth, so training requires networks of machines. It introduces the hardware hierarchy — fast NVLink within a machine's eight GPUs, slower InfiniBand across machines, and an even slower tier beyond roughly 256 GPUs — which shapes every later design decision.

Collective communication and the all-reduce identity (4:09)

It reviews all-reduce, broadcast, reduce, all-gather and reduce-scatter, then establishes that all-reduce is equivalent in bandwidth to a reduce-scatter followed by an all-gather. It briefly contrasts GPU all-to-all networking with TPU toroidal mesh networking, noting collectives can be implemented efficiently on either.

Data parallelism and ZeRO/FSDP (12:14)

Naive data parallelism replicates parameters, gradients and optimizer state on every GPU, which scales compute well but wastes memory badly (roughly 16x parameter memory once Adam's moment estimates are counted). The lecture builds up ZeRO stage by stage: stage 1 shards optimizer state, stage 2 additionally shards gradients by sending them to their owning GPU as they are computed during the backward pass, and stage 3 (equivalent to FSDP) shards parameters too, fetching and freeing them on demand while overlapping communication with computation to keep overhead low.

Pipeline parallelism (45:34)

Pipeline parallelism cuts the model by layer, with each GPU passing activations to the next. The naive version leaves GPUs mostly idle (a large bubble); splitting work into micro-batches overlaps stages and shrinks the bubble in proportion to batch size. The lecture also describes zero-bubble/dual-pipe scheduling, which separates the backward computation for activations from the backward computation for weight gradients so the latter can fill otherwise idle time, at the cost of substantial implementation complexity.

Tensor parallelism (55:38)

Tensor parallelism splits matrix multiplies into submatrices computed on different devices, synchronized with all-reduces in both forward and backward passes. It avoids the bubble problem and does not consume batch size, but needs high-bandwidth interconnects, so it is applied mainly within a single machine's roughly eight GPUs; throughput drops sharply beyond that.

Sequence parallelism and activation memory (1:05:44)

Tensor and pipeline parallelism reduce most memory linearly but leave point-wise operations like layer norm and dropout unsharded. Sequence parallelism splits these along the sequence dimension, using all-gathers and reduce-scatters to synchronize, closing the remaining gap in activation memory.

Putting it together, with examples (1:16:49)

The lecture gives a rule of thumb: use tensor parallelism up to a machine's GPU count, then FSDP or pipeline parallelism across machines until the model fits in memory, then data parallelism (with gradient accumulation if needed) to use remaining GPUs. It closes with real examples from Megatron-LM, OLMo, DeepSeek, Llama 3 and Gemma 2, including Llama 3's reported hardware failure rates during training.

Before you watch

• Be comfortable with single-GPU throughput optimization and memory accounting for parameters, gradients and optimizer state, covered in the previous lecture.
• Know the basics of the Adam optimizer's first and second moment estimates.
• Familiarity with flash attention's tiling idea is helpful for the ring attention mention near the end.

Check your understanding

1. Why does an all-reduce cost the same bandwidth as a reduce-scatter followed by an all-gather, and why does that matter for ZeRO?
2. Walk through what changes at each ZeRO stage (1, 2, 3) in terms of what is sharded and what is communicated.
3. Why is tensor parallelism usually limited to the GPUs within a single machine, while pipeline parallelism is often used across machines?
4. Why is batch size described as a limited resource, and what are two different ways of 'spending' it?
5. What memory does tensor parallelism fail to reduce, and how does sequence parallelism address it?