Pipeline Parallelism from First Principles
Splitting by depth
Data parallelism replicates the model and splits the batch. Pipeline parallelism does the opposite: it splits the model and passes the batch through. With $P$ stages, stage $p$ owns a contiguous run of $L/P$ layers.
What it costs to communicate
Only the activations crossing stage boundaries, a tensor of shape $(B,S,d)$, sent point-to-point to exactly one neighbour. There is no collective and no all-reduce. For a 4096-wide model this is small compared to sending gradients.
What it buys
Each device stores $1/P$ of the parameters, gradients and optimizer state — the same factor as ZeRO-3, but with far less traffic. Pipelining is the cheapest way to spread a model across nodes over a slow network.
The bubble
Run one batch through a 4-stage pipeline and watch the devices. Stage 0 computes, then stage 1, then 2, then 3; the backward pass returns in reverse. At every instant exactly one device is busy.
Four devices to get one device's throughput is worse than useless. The fix is to break the minibatch into $m$ microbatches and push them in one after another, so that while stage 1 works on microbatch 1, stage 0 has already started microbatch 2.
Count it. Let one microbatch's forward-plus-backward cost on one stage be $t$. A perfect pipeline would take $m\,t$. The real pipeline needs $P-1$ extra slots to fill and drain:
Two readings of that formula matter. It falls as $m$ grows, and it rises with $P$. Setting it below a target $\epsilon$ requires
so a 10% bubble needs $m\ge 9(P-1)$: 28 microbatches for 4 stages, 64 for 8, 136 for 16.
| Stages $P$ | $m=4$ | $m=8$ | $m=16$ | $m=32$ | $m=64$ |
|---|---|---|---|---|---|
| 2 | 20.0% | 11.1% | 5.9% | 3.0% | 1.5% |
| 4 | 42.9% | 27.3% | 15.8% | 8.6% | 4.5% |
| 8 | 63.6% | 46.7% | 30.4% | 17.9% | 9.9% |
| 16 | 78.9% | 65.2% | 48.4% | 31.9% | 19.0% |
The memory that pipelining creates
Microbatching is not free, and the price is activation memory. A microbatch's activations must be kept from the moment its forward pass runs until its backward pass consumes them. In a plain all-forward-then-all-backward schedule, stage 0 runs every forward before any backward, so it holds activations for all $m$ microbatches at once.
The imbalance is also spatial. Stage 0 holds the most activations and stage $P-1$ the fewest, because the last stage's microbatches are consumed almost immediately. Memory is therefore lopsided across the pipeline even when compute is perfectly balanced.
Balancing the stages
The pipeline runs at the speed of its slowest stage, so an uneven split wastes time on every microbatch, not just once. Equal layer counts are the obvious split and usually the wrong one:
| Source of imbalance | Effect | Mitigation |
|---|---|---|
| Embedding table on stage 0 | Large parameter and gradient memory, but almost no compute. | Give stage 0 extra transformer blocks to compensate. |
| Output head and loss on the last stage | A $V\times d$ matmul plus a softmax over the vocabulary is genuinely expensive, and the logits tensor is huge. | Give the last stage fewer blocks; compute the loss in chunks. |
| Tied embeddings across stages 0 and $P-1$ | One tensor lives on two devices and receives gradients on both. | An explicit all-reduce between those two ranks before the step. |
| Heterogeneous devices or interconnects | The slowest link throttles everything. | Profile, then assign layers by measured time rather than by count. |
What the implementation must get right
Send and receive in a deadlock-free order. If every stage posts a blocking receive before its send, the pipeline deadlocks instantly. Either alternate the order by parity of the rank, or use non-blocking primitives and wait afterwards.
Shapes must be known in advance. The receiver has to allocate before the data arrives. Variable sequence lengths mean either padding to a fixed shape or sending metadata first.
Accumulate gradients across microbatches. Each microbatch produces a partial gradient; the optimizer must step once, after all $m$ have been accumulated. This is the time-domain version of the same summation argument from the autograd note.
Scale the loss by $1/m$. Otherwise the accumulated gradient is $m$ times too large — the same trap as gradient accumulation without pipelining.
Pipeline versus tensor parallelism
| Pipeline | Tensor | |
|---|---|---|
| Split along | layers (depth) | inside each matmul (width) |
| Communication | point-to-point, once per stage boundary | collective, several times per block |
| Volume per step | low | high |
| Tolerates slow links | yes — use across nodes | no — keep inside a node |
| Main cost | the bubble, and activation memory | bandwidth, and synchronization |
They are complementary, which is why production runs combine them: tensor parallel within a node over the fast fabric, pipeline parallel across nodes over the slow one, data parallel on top. The device-mesh note assembles exactly that.
Pipelining gives you a $1/P$ memory reduction for almost no bandwidth, at the cost of a bubble of $(P-1)/(m+P-1)$. Microbatching shrinks the bubble in proportion to $m$ and grows activation memory in proportion to $m$. Resolving that trade is what pipeline schedules are for.
Check yourself
- Derive the bubble fraction $(P-1)/(m+P-1)$ from the fill-and-drain argument. derivation
- Invert it to find the $m$ needed for a 5% bubble at $P=8$. calculation
- Explain why stage 0 holds more activations than the last stage. reasoning
- Give a layer assignment for $L=32$, $P=4$ that accounts for the embedding and the output head. design
- Describe a send/receive ordering that deadlocks, and one that does not. debugging