Pipeline Parallel and Bubble Analysis

A 70B-parameter model in fp16 needs 140 GB for parameters alone, which no consumer GPU holds. ZeRO-3 shards the parameters across ranks, but each rank still has to allgather the full layer for each forward step, paying log(N) hops per layer. Pipeline parallel takes a different route. Cut the model into N stages and put one stage on each rank. Layer 1's forward finishes on rank 0 and hands the activation tensor to rank 1, rank 1 runs layer 2 and hands off to rank 2, and so on. Backward flows in reverse. Memory drops linearly because each rank holds one stage. Compute is sequential, and that is where the bubble comes from.

The bubble is idle time at the start of the pipeline (waiting for the first microbatch to reach the last stage) and at the end (waiting for the last microbatch to drain back through). With M microbatches and N stages, the per-stage bubble fraction is (N-1)/(M+N-1). At M=8, N=4 that is 27%. At M=64, N=4 it is 4.5%. More microbatches per step shrink the bubble, which means smaller per-microbatch batch sizes, and that constraint drives microbatch design.

GPipe schedule

GPipe fills the pipeline forward with all M microbatches before starting any backward pass, then drains backward in reverse. Activations from each microbatch are held until its backward runs, so memory grows linearly with M. Forward takes M+N-1 cycles and backward takes another M+N-1. Per-stage useful work is 2M cycles and the per-stage bubble is 2(N-1) cycles. When each forward and backward takes one unit of time, the bubble fraction is (N-1)/(M+N-1). Picking M much greater than N hides the bubble.

1F1B schedule

1F1B interleaves the two passes: as soon as a microbatch's forward reaches the last stage, its backward starts and streams back. Each stage alternates one forward and one backward. The bubble is still N-1, but activation memory is bounded by the pipeline depth instead of the microbatch count. Megatron and PipeDream both use 1F1B. GPipe comes first here because it is simpler, and 1F1B is left as an exercise.

Why equal compute per stage matters

Say stage 0 takes 50 ms and stage 1 takes 100 ms. Each cycle is gated on stage 1, and the other stages sit idle for 50 ms per cycle waiting for it to release. Equal parameter count is the wrong axis. A transformer's compute is dominated by attention plus MLP in each layer, while embedding layers have many parameters and little compute. Assign stages to equalise FLOPs per stage, not weights per stage.

Microbatch versus batch

A pipeline runs M microbatches of size B each, so the effective batch size is MB. The gradient at the end of a pipeline step is the gradient over those MB examples combined. The bubble fraction depends on M, while the optimiser sees M*B. Tuning M trades bubble (lower with high M) against per-microbatch memory (under GPipe, activation memory grows with M).

code/main.py implements:

• StageMLP: a small nn.Module that holds one stage's two linear layers and exposes forward(x).
• gpipe_schedule(num_stages, num_microbatches): lays out the GPipe schedule as (cycle, stage, microbatch, phase) tuples in integer cycles, with every forward and backward costing one cycle. (FORWARD_UNITS and BACKWARD_UNITS are defined at the top of the file but the schedule does not use them.)
• render_gantt(...): draws that schedule as a stage-by-cycle grid with microbatch labels such as F3 and B3.
• bubble_fraction(num_stages, num_microbatches): the closed form (N-1)/(M+N-1).
• measure_bubble(...): counts the idle slots in the schedule, which is the measured fraction.
• run_pipeline(...): a second demo, a real two-rank pipeline over torch.distributed with the gloo backend. Rank 0 runs stage 0 and sends the activation to rank 1, rank 1 runs stage 1, computes the loss and sends the gradient back, and both ranks train a small MLP for a few steps.

Run it:

python3 code/main.py

The output is a 4-stage, 8-microbatch Gantt grid, the measured bubble percentage next to the closed-form prediction, a table of bubble versus M, and then a short summary of the two-rank gloo run (microbatch losses seen by rank 1 and the final parameter norms).

Three patterns harden pipeline parallel enough to ship.

Activation checkpointing pairs with pipeline. With M microbatches in flight on GPipe, activation memory is M times one microbatch. Activation checkpointing recomputes the forward at backward time, trading compute for memory, and the combination is what makes pipeline tractable for long sequences.

Stage balance is measured, not assumed. Production teams run a profiling pass that measures actual per-layer compute (FLOPs and wall-clock) on the target hardware, then partition by that measurement. The Megatron-LM --num-layers-per-stage flag accepts a list, which allows uneven layer counts when stages have different per-layer cost.

The send-recv schedule has to avoid deadlock. If each stage sends before it receives, the pipeline deadlocks on the wire. The standard fix is to interleave: even-rank stages send first and then recv, odd-rank stages recv first and then send. code/main.py only has two ranks, so the pattern shows up in its smallest form: rank 0 sends the activation and then receives the gradient, while rank 1 receives first and then sends.

Production patterns:

• Megatron-LM. The reference for pipeline parallel at scale. It uses 1F1B and supports tensor, pipeline, and data parallel combined.
• DeepSpeed Pipeline. Integrates with ZeRO. ZeRO-1 plus pipeline is a common combination for the largest open models.
• PyTorch Pipe. The PyTorch-native pipeline wrapper, built on torch.distributed.pipeline.sync.Pipe.

Lesson 80 stores the per-stage parameter shards in the sharded checkpoint. Lesson 81 composes DDP, ZeRO, and pipeline in the end-to-end demo, in spirit: the demo keeps the pipeline simulated to keep runtime short.

1. Implement 1F1B and verify the bubble fraction matches GPipe but activation memory is bounded.
2. Profile real per-stage time on a deeper model and rebalance stages by measured wall-clock.
3. Add gradient accumulation across pipeline microbatches and check the gradient equals the gradient of the equivalent full-batch forward.
4. Pair the pipeline with activation checkpointing and measure the memory drop versus compute cost.
5. Combine pipeline with DDP (each pipeline rank is replicated across a data-parallel group) and reason through the 2D schedule.

• Huang et al, GPipe: Efficient Training of Giant Neural Networks
• Narayanan et al, PipeDream: Generalized Pipeline Parallelism for DNN Training
• Megatron-LM pipeline parallel docs
• Phase 19 Lesson 76 - the send/recv primitives the schedule uses
• Phase 19 Lesson 78 - ZeRO is orthogonal to pipeline and often combined