Deep Learning · Efficient Training
Pipeline Parallelism: Filling the Bubble and Paying for It
Micro-batch pipelining lets layers be split across devices, but the schedule leaves idle slots. How bubble fraction scales with depth and why interleaved schedules help.
Splitting a model’s layers across devices creates a sequential dependency: layer 40 cannot start until layer 39 has produced its output, and layer 39 lives on another device. A single forward pass therefore uses one device at a time and leaves the rest idle, which is worse than not splitting at all.
Micro-batch pipelining fills the idle time by interleaving independent work. The batch is divided into micro-batches, and while device 2 processes micro-batch 1, device 1 processes micro-batch 2. The steady state keeps every device busy, and the cost is the ramp-up and drain periods at the beginning and end of the step where the pipeline is not yet full.
The schedule
With $P$ pipeline stages and $M$ micro-batches, the standard schedule issues forward passes in an order where stage $s$ handles micro-batch $m$ at time slot $s + m$. Backward passes follow in reverse, once the corresponding forward has reached the last stage.
The first $P-1$ time slots are the fill: stage 1 runs micro-batch 1 alone, then stages 1 and 2 run micro-batches 2 and 1, and so on. The last $P-1$ slots are the drain, where stages finish their remaining micro-batches and idle. In between, all $P$ stages are active on every slot.
Total slots per step are $M + P - 1$ forward slots and the same for backward, while the useful work is $M$ slots per stage. The idle fraction is
$$\text{bubble} = \frac{P - 1}{M + P - 1}$$
which is the fraction of time each stage spends waiting. For $P = 8$ and $M = 8$, the bubble is $7/15 \approx 47$ percent, which is unacceptable. For $M = 64$, it is $7/71 \approx 10$ percent. The fraction falls as $1/M$ for large $M$, so the remedy is more micro-batches, and more micro-batches means a smaller micro-batch for a fixed total batch, which brings back the arithmetic-intensity problem from the compute accounting.
There is a second constraint: each micro-batch’s activations must remain resident from its forward pass until its backward pass, which is $P$ stages later. Activation memory per stage therefore scales with $M$, not with the micro-batch size alone. Doubling $M$ to halve the bubble doubles the stored activations. This is the central tension of pipeline scheduling, and it is why checkpointing inside a pipeline stage is common: it breaks the $M$-proportional term back down.
The forward schedule for four stages and four micro-batches makes the idle slots concrete; each column is one time slot and a dot marks a stage with nothing to do.
slot: 1 2 3 4 5 6 7
stage 1: m1 m2 m3 m4 . . .
stage 2: . m1 m2 m3 m4 . .
stage 3: . . m1 m2 m3 m4 .
stage 4: . . . m1 m2 m3 m4
Four micro-batches over four stages gives a bubble of $3/7$, and the drain at the end repeats the same pattern in reverse for the backward passes.
Interleaving to shrink the bubble
The bubble formula assumes each stage owns one contiguous block of layers. An alternative assigns each device several non-contiguous chunks and processes them in a round-robin order, so that a device’s second chunk can run while other stages work on the first.
With $V$ chunks per device, the bubble becomes approximately $\frac{P-1}{V \cdot M + P - 1}$. Doubling the chunks per device roughly halves the idle fraction for the same $M$. The cost is more communication events: each chunk boundary is a pipeline transfer, so the number of sends and receives grows with $V$, and each has latency.
The scheduling complexity also grows. The interleaved order must respect both the layer dependency and the chunk assignment, and the backward order must mirror the forward order exactly or the gradients are applied to the wrong activations. Implementations that support this expose it as a number of chunks or virtual stages, and the useful range is small: past a handful of chunks, the added latency and bookkeeping outweigh the bubble reduction.
The memory arithmetic per stage
Pipeline splitting reduces parameter memory per device by $P$, because each device holds $1/P$ of the layers. It does not reduce the activation memory per layer, and it multiplies the number of resident micro-batches by $M$.
The optimiser state follows the parameters, so it also divides by $P$. For a model with $N$ parameters and Adam in mixed precision at 18 bytes per parameter, a stage holds $18N/P$ bytes of state, plus its share of activations.
The balancing problem is that layers are not equal in cost. A transformer’s embedding and its final projection are cheap in compute and large in parameters; the middle layers are uniform. A split that gives each stage the same number of layers can give them very different memory and compute requirements, and the pipeline runs at the speed of its slowest stage. Balanced partitioning by measured per-layer cost, rather than by layer count, is the difference between a pipeline that reaches its predicted throughput and one that does not. Since a forward and a backward pass through the same layer cost different amounts, the balance must be struck on the combined cost.
What goes across the stage boundary
Only the activation tensor at each boundary crosses devices: for a transformer, one tensor of shape (micro-batch, sequence, hidden) per boundary. That is small compared with the parameter traffic of a sharded optimiser, which is why pipelines scale across nodes more gracefully than sharding does.
The size still matters for latency. A boundary tensor of $b \times S \times H$ elements in bf16 is $2bSH$ bytes. For $b = 1$, $S = 2048$, $H = 4096$, that is 16 MB, transferred twice per micro-batch (forward and backward). The transfer time on a link with tens of gigabytes per second is on the order of a millisecond, comparable to the compute time of a single micro-batch in a deep model. At that point the pipeline is latency-bound at every boundary, and the schedule’s idle slots are no longer the only source of loss.
Gradient accumulation across micro-batches interacts here. If accumulation is used inside a stage, the boundary transfer happens once per accumulation window rather than once per micro-batch, which reduces boundary traffic but lengthens the pipeline’s effective step. The two mechanisms are solving the same memory problem from different directions and must be co-designed rather than layered.
Backward pass ordering and weight updates
The straightforward schedule performs all forward passes, then all backward passes, then one optimiser step. That is what the bubble formula describes, and it requires all $M$ micro-batches’ activations to be resident simultaneously.
An alternative interleaves forward and backward more aggressively: as soon as the last stage finishes the forward pass for micro-batch $m$, the backward pass for micro-batch $m$ can begin at the last stage, freeing that micro-batch’s activations earlier. This reduces peak activation memory at the cost of a more complex schedule and, in some formulations, of performing the weight update once per micro-batch rather than once per batch.
Updating weights per micro-batch changes the optimisation: it is equivalent to a smaller effective batch with more frequent steps, which is a different training procedure, not merely a different schedule. The distinction between schedules that preserve batch semantics and those that do not is the same distinction that appears in gradient accumulation, and it is easy to lose when comparing pipeline implementations by throughput alone.
Where pipelines fit
A pipeline is the right tool when the model does not fit on one device and the interconnect between devices is slower than the intra-device fabric. It divides parameters and optimiser state with modest communication per boundary, and its cost is a scheduling bubble that shrinks with the number of micro-batches.
It is the wrong tool when the bottleneck is activations at long sequence length, because the pipeline does not reduce per-layer activation size and it multiplies resident micro-batches. It is also the wrong tool when the model fits with sharding alone, since sharding adds no bubble.
Real large-scale configurations combine the two: pipeline stages across nodes, and within each stage a sharded data-parallel group, sometimes with tensor parallelism splitting individual layers further. Choosing the sizes of those three dimensions under a fixed device count and memory budget is the composition problem treated in planning a parallelism layout.
References
- GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism — Huang et al., NeurIPS 2019
- PipeDream: Fast and Efficient Pipeline Parallel DNN Training — Harlap et al., 2018
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al., 2019