Prady Prakash

Module 9

Pipeline Parallelism

Tensor parallelism ran out of room at the edge of the node — its communication is too heavy for InfiniBand. So to split a model across many nodes you need a strategy whose communication is light enough to survive the slow links. Pipeline parallelism is that strategy, and it communicates the least of any model-splitting approach: it only passes activations at layer boundaries, once per microbatch, point to point.

The catch is a new failure mode that none of the previous chapters had — GPUs sitting idle not because they're waiting on memory or network, but because they have nothing to do yet. Managing that idle time, the "bubble," is the entire subject of this chapter, and it's the one place where the schedule — the pure choreography of which GPU does what when — is the whole game.

The naive version, and why one GPU works at a time

Split the model by layers: stage 0 holds layers 0–7, stage 1 holds 8–15, and so on across dd stages. A batch flows through stage 0, whose output feeds stage 1, and so forth; gradients flow back in reverse.

Implemented naively, this is worse than a single GPU. Stage 1 can't start until stage 0 finishes; stage 2 waits on stage 1. At any instant exactly one stage is busy and d1d-1 are idle. You've bought dd GPUs to get 1/d1/d the utilization — the layers fit now, but you're wasting almost the entire cluster. This is real model parallelism and it's why "just split the layers" isn't the answer.

Microbatching, and the bubble

The fix is to keep every stage fed. Split each minibatch into mm microbatches and pipeline them: as soon as stage 0 finishes microbatch 0, it hands it to stage 1 and immediately starts microbatch 1. Once the pipeline fills, every stage is working on a different microbatch simultaneously. This is GPipe.Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019.

But the fill and drain aren't free. At the start, stages 1d11 \ldots d-1 wait for the first microbatch to reach them; at the end, stages 0d20 \ldots d-2 sit idle as the last microbatch drains out. That idle triangle is the bubble, and its size has a clean closed form.

The pipeline takes d1d-1 steps to fill and d1d-1 to drain, over a useful span of mm steps. So:

bubble fraction=d1m+d1\text{bubble fraction} = \frac{d-1}{m + d - 1}

This is the equation that governs pipeline parallelism, and it says one thing: make mm large relative to dd. At d=8d = 8, m=8m = 8 you're idle 47% of the time — catastrophic. At d=8d = 8, m=64m = 64 it's 10%. The lever is the ratio.

Watch it directly — start on GPipe, then push the microbatch slider and watch the bubble shrink:

stage 0
0
1
2
3
4
5
6
7
0
1
2
3
4
5
6
7
stage 1
0
1
2
3
4
5
6
7
0
1
2
3
4
5
6
7
stage 2
0
1
2
3
4
5
6
7
0
1
2
3
4
5
6
7
stage 3
0
1
2
3
4
5
6
7
0
1
2
3
4
5
6
7
ForwardBackward (∂input)idle (bubble)
Pipeline depth (stages)4
Microbatches8

27.3% of the pipeline is idle bubble. GPipe's bubble is (d−1)/(m+d−1) = (4−1)/(8+4−1). More microbatches shrink it, but every in-flight forward's activations are held until its backward runs — so the memory cost grows with m.

GPipe versus 1F1B: a memory story, not a bubble story

Here's the subtlety that trips people up, and it's worth being precise about because the two schedules have identical bubble fractions.

GPipe runs all mm forward passes, then all mm backward passes. Simple, but it means every microbatch's activations must be kept alive from its forward until its backward — so peak activation memory scales with mm, the very thing you wanted to make large to shrink the bubble. GPipe forces you to choose between a small bubble and fitting in memory.

1F1B (one-forward-one-backward) breaks the tension.Narayanan et al., PipeDream: Generalized Pipeline Parallelism for DNN Training, SOSP 2019, introduced the 1F1B schedule; Megatron-LM adopted the memory-efficient synchronous variant. After a warmup of dd forwards to fill the pipe, each stage alternates: one forward, one backward, steady state. The instant a microbatch's backward completes, its activations are freed. Peak activation memory now scales with dd (the pipeline depth) instead of mm — and dd is small and fixed.

Switch the visualizer to 1F1B: the bubble is the same size, but the backwards now interleave with forwards from the start instead of bunching at the end. That interleaving is the whole point. Same bubble, far less memory, which means you can afford the large mm that shrinks the bubble. It's the default in every serious training framework, and the reason is memory, not idle time.

Interleaving and the weight-staleness trap

You can shrink the bubble further by assigning each GPU multiple non-contiguous stages — GPU 0 gets layers 0–3 and 16–19, so a microbatch visits it twice. This "interleaved" or "virtual pipeline" schedule cuts the bubble by the number of chunks per device, at the cost of proportionally more point-to-point communication. Megatron uses it heavily.

But the moment you interleave forwards and backwards from different microbatches, a correctness question appears: a microbatch's forward and backward passes might see different versions of the weights if an optimizer step landed in between. That's weight staleness, and it silently degrades learning if you ignore it. The asynchronous PipeDream variants had to solve it explicitly:

  • Weight stashing — each stage keeps several weight versions and uses the matching one for each microbatch's backward. Correct, but memory-hungry.
  • PipeDream-2BW (double-buffered weights) — keep only two versions, generating a new one every kk microbatches. Bounded memory, bounded staleness.
  • PipeDream-flush (= the synchronous 1F1B everyone actually uses) — periodically flush the pipeline and do a global synchronized step, so there is no staleness. You pay a bubble at each flush; in exchange the math is exactly single-GPU SGD.

Modern synchronous training takes the flush. The staleness-tolerating variants are clever and mostly historical for dense LLM training — worth knowing because the vocabulary still shows up, not because you'll reach for them.

Zero-bubble: splitting the backward pass

The bubble seems fundamental — filling and draining a pipe takes time. But there's a trick, and it comes from looking harder at what "backward" means.

A backward pass through a layer computes two different gradients: the gradient with respect to the layer's input (call it B, needed immediately by the previous stage to continue its backward) and the gradient with respect to the layer's weights (call it W, needed only at the optimizer step, which doesn't happen until the very end). Standard schedules fuse B and W into one op. But only B is on the critical path.

Zero-bubble pipeline parallelism splits them.Qi et al., Zero Bubble Pipeline Parallelism, ICLR 2024. The insight that B and W have different scheduling constraints is simple in hindsight and was genuinely new. The B passes keep the dependency chain moving; the W passes are free-floating and get slotted into what would otherwise be bubble. Since W has no downstream dependency, it's the perfect filler. Their ZB-H1 schedule reduces the bubble substantially; ZB-H2, by also reshaping the warmup and reordering the tail, drives it to essentially zero — turning the schedule's trapezoid into a gap-free parallelogram.

Flip the visualizer to zero-bubble and watch the W cells (a third color) migrate into the gaps the other schedules leave empty.

DeepSeek-V3's DualPipe pushes the idea further still, running bidirectional pipelines that feed microbatches from both ends and overlapping the forward/backward compute of one chunk with the communication of another — much of the pipeline's communication disappears behind computation entirely.DeepSeek-AI, DeepSeek-V3 Technical Report, 2024.

What Llama 3 actually changed

The clean (d1)/(m+d1)(d-1)/(m+d-1) story assumes every stage takes equal time and holds equal memory. Real transformers violate both assumptions at the ends, and Llama 3's pipeline work was largely about fixing the resulting imbalances.Grattafiori et al., The Llama 3 Herd of Models, 2024, §3.3.2 covers the pipeline modifications directly.

Two imbalances specifically:

  • Memory imbalance. The first stage holds the token embedding; the last holds the output projection and the loss computation over a 128k-token vocabulary — which, from chapter 2, is a genuinely large activation. The last stage also holds the most in-flight microbatches under 1F1B. So the end stages carry more memory than the middle ones, and the whole pipeline is capped by whichever stage OOMs first.
  • Computation imbalance. That same vocabulary projection and loss make the last stage's compute heavier, so it becomes the latency bottleneck every other stage waits on.

The fix Llama 3 used is pleasingly blunt: rebalance the layers. Drop one transformer layer from the first stage and one from the last, so the first stage's lighter transformer load makes room for the embedding, and the last stage's makes room for the output head and loss. They also modified the schedule to run an arbitrary number of microbatches per batch, relaxing a constraint the standard 1F1B implementation imposes, to keep the bubble small across the batch sizes the run actually used.

What makes this imbalance nasty in practice is how it presents. It rarely announces itself as "the pipeline is imbalanced." It shows up as the last stage OOMing while nvidia-smi on the middle stages shows gigabytes free — memory you can't use, because it's on the wrong device. Or it shows up as a throughput number that refuses to improve no matter how you tune the batch, because the last stage's heavier compute is silently setting the pace for everyone. The tell is a per-stage memory or timing trace that isn't flat: one or both ends stick up above the middle. Once you're looking at that trace the fix is obvious, and "drop a layer from each end" is exactly the kind of blunt rebalancing that flat trace asks for. Getting there is mostly about knowing to plot per-stage curves at all rather than staring at an aggregate that averages the imbalance away.

The general lesson generalizes past Llama: the analytical bubble formula is the start of pipeline tuning, not the end. Real pipelines are bottlenecked by their least-balanced stage, and the embedding/loss asymmetry at the ends is the usual culprit. Getting a pipeline to its theoretical efficiency is mostly the work of flattening those two curves.