Pipeline parallelism schedules

Background

Pipeline parallelism splits a model's layers into K sequential stages, assigns each stage to its own process, and streams microbatches through them, so that a model whose parameters and activations exceed a single device's memory can still be trained. The cost is the pipeline bubble. At the start of every batch, downstream stages sit idle until the first microbatch reaches them, and at the end, upstream stages sit idle until the last gradients come back. These two phases are the pipeline fill and drain, and together they cost on the order of K ticks per batch. The schedule decides how much of that time is recovered and at what cost in memory and communication.

The three standard schedules make different trades:

• GPipe runs all M forward passes, then all M backward passes. The bubble shrinks relative to useful work as M grows, but every stage must keep the activations of all M microbatches until their backward passes arrive, so activation memory grows with M.
• 1F1B alternates one forward and one backward pass after a short warm-up. Backward passes start before all forward passes have finished, which removes the long idle gap between the two phases. A stage never holds more than K microbatches of activations, so memory is bounded by pipeline depth rather than by M.
• Interleaved 1F1B splits each stage into v smaller virtual stages. Fill and drain shorten by roughly a factor of v, in exchange for v times more hand-offs between processes.

The diagram below shows the first two on four stages and four microbatches. The bubble is the same size in both; what differs is how long activations stay resident.

Benchmark

I built a runner on PyTorch's pipelining API that sweeps schedule, model depth and width, process count, microbatch count and interleaving factor, records tokens per second for each configuration, and plots speedup over GPipe and scaling efficiency.

The runs use CPU processes with the gloo backend. Without GPU kernels or NCCL involved, the measurement isolates the schedules' own overheads, and it puts them in a setting where communication is expensive relative to compute, so their differences are easy to see. The chart below reports throughput for each schedule on a 4-layer transformer at two widths (4 and 8 attention heads) and at two and four processes.

• On the narrower model at two processes, 1F1B reached 1.20x GPipe's throughput and interleaving reached 1.13x. Stages are cheap, so removing the middle bubble pays off and the extra coordination costs little.
• On the wider model at two processes the order flipped: interleaving edged ahead at 1.03x and 1F1B fell slightly below GPipe. With more compute per stage, fill and drain become a larger share of each step, and interleaving is the schedule that shortens them.
• At four processes all three schedules finished within a few percent of each other, and scaling efficiency dropped from 50 to 60% to about 25%. Synchronization overhead from barriers, context switches and tensor hand-offs grew faster than useful work per stage. GPipe makes the fewest of them, so it was competitive or best.

Takeaway

Which schedule is fastest depends on the ratio of compute per stage to communication cost. Overlapping schedules recover idle time but add synchronization, and when communication dominates, as it does on CPU at four processes, the added synchronization costs more than the bubble it removes. On GPUs the ratio is reversed. NCCL makes communication cheap and each stage carries far more compute, so the overlapping schedules come out ahead in most configurations. That is the setting they were designed for.

The memory result does not depend on hardware. 1F1B's activation footprint is set by pipeline depth, not by the number of microbatches in flight, so a deep pipeline can run enough microbatches to hide its bubble without running out of memory. This is why 1F1B is the standard schedule in large-scale training frameworks, whatever throughput margin it shows on a particular machine.