Volume II · Splitting the Work

Stages, Bubbles, and Redundant Memory

Tensor parallelism cuts inside a layer. This lecture covers the other way to slice a model that does not fit on one GPU: cut between layers instead, and give each GPU a contiguous stretch of the model to own outright. That trade buys back memory at the cost of idle time — and a second, orthogonal idea, ZeRO, buys back a different kind of memory that was never doing useful work in the first place.

Pipeline parallelism: splitting by layer

Where tensor parallelism splits the weight matrix inside one layer across GPUs, pipeline parallelism splits the model by layer: GPU 0 owns layers 1 through k, GPU 1 owns layers k+1 through 2k, and so on, with activations physically passed from one GPU to the next as data flows through the network. Huang et al.'s GPipe (arXiv:1811.06965) is the paper that made this practical at scale: it splits each training batch into smaller micro-batches and pipelines them synchronously through the stages, which lets near-linear speedup emerge from adding more stages — but it pays for that synchrony with "bubble" time, the idle GPU cycles that occur while the pipeline is filling up at the start of a batch and draining at the end, because early stages have nothing left to do once their micro-batches have all moved downstream and late stages have nothing to do until the first micro-batch arrives.

Harlap et al.'s PipeDream (arXiv:1806.03377) attacks exactly this idle time with an asynchronous schedule: a "one-forward-one-backward" pattern that keeps every stage doing useful work far more of the time than GPipe's synchronous flush allows, at a real cost — because the schedule is asynchronous, some weight updates are computed using slightly stale versions of the weights rather than the very latest ones. The trade is explicit: GPipe accepts idle bubble time to keep every update exactly correct; PipeDream accepts slightly stale updates to keep GPUs almost always busy.

ZeRO: memory that was redundant, not shared

Data parallelism (Lecture 3) has a quieter cost than the bubble time above: every GPU in a data-parallel group holds a full copy of the optimizer state, the gradients, and the parameters — not because each GPU needs a private copy to do useful work, but simply because that is what naive replication does. Rajbhandari et al.'s ZeRO (arXiv:1910.02054, Microsoft, 2019/2020) names this plainly: that replicated memory is redundant, not shared in any useful sense, and it can be partitioned across the data-parallel GPUs instead of copied onto every one of them.

ZeRO does this in three increasingly aggressive stages. Stage 1 partitions only the optimizer state across GPUs — each GPU keeps its full copy of gradients and parameters, but the (often large, e.g. Adam's two moment buffers per parameter) optimizer state lives on only one GPU per shard. Stage 2 additionally partitions the gradients. Stage 3 goes further still and partitions the parameters themselves, so no single GPU ever holds a complete copy of the model's weights at rest — only reconstructing the piece it needs, communicating to fetch the rest, when a computation actually requires it. The paper's own claim is that this could in principle scale model size roughly linearly with the number of GPUs, rather than being capped by whatever fits in one GPU's memory, at a communication overhead the paper argues stays low relative to the memory it frees. PyTorch's Fully Sharded Data Parallel (FSDP) is the widely-used implementation of this same idea, most directly corresponding to ZeRO stage 3.