03 — Pipeline Parallelism

Papers: GPipe (Huang et al. 2018, arxiv:2507.20534), PipeDream (Narayanan et al. 2019), Megatron-LM 2 (arxiv:2310.01889) for interleaved 1F1B, Zero Bubble Pipeline Parallelism (Qi et al. 2023, arxiv:2310.01889), PipeOffload (Wan et al. 2025) for the memory story.

The idea in one sentence

Instead of splitting a matmul across GPUs (TP) or a batch across GPUs (DP), split the layers themselves into consecutive stages. GPU-0 owns layers 0–19, GPU-1 owns 20–39, etc. A batch enters at stage 0 and flows through.

Why PP exists

  • The model doesn’t fit on one node’s NVLink domain (past ~350B params in bf16), and TP is capped by NVLink domain size (~8 on H100 SXM).

  • Cross-node communication needs to be minimal-and-structured. PP sends only the activation (batch, seq, hidden) between stages, not a per-block collective. This is the only parallelism that survives InfiniBand gracefully.

  • Memory sharding by layers is coarse-grained but very effective: a 405B on 8 nodes with PP=8, TP=8 = 64-GPU 4D config is the canonical frontier setup.

The bubble problem

Naive PP: stage 0 processes the whole batch, hands off, stage 1 processes it, etc. Only one stage is active at a time. If you have p stages, p-1 stages idle per timestep. This is the pipeline bubble.

Bubble fraction (naive) = (p-1) / p. At p=8, 87.5% of compute is wasted. Unacceptable.

GPipe (2018) — microbatching

Split the mini-batch into m micro-batches. Feed them through the pipeline back-to-back. Now the pipeline fills, works at full capacity, drains.

Bubble fraction = (p-1) / (m + p - 1).

At p=8, m=32: bubble = 7/39 = 18%. Better.

Cost: you must stash activations for all m micro-batches for the backward pass. Activation memory blows up linearly with m. GPipe adds activation checkpointing/recompute to control this, but recompute costs ~33% throughput.

1F1B (One-Forward-One-Backward) — PipeDream style

Interleave forward and backward passes: as soon as a stage finishes one micro-batch’s forward and receives that micro-batch’s backward from the downstream stage, do the backward. This is the standard schedule in Megatron.

Bubble is the same size but peak activation memory drops dramatically because a stage stashes only p micro-batches (its “in-flight” ones), not all m.

This is the default PP schedule everywhere. GPipe is legacy.

Interleaved 1F1B (Megatron-2)

Instead of stage 0 owning layers 0–19 contiguously, give it layers {0,1,10,11} (two virtual chunks). Now the pipeline has p × v virtual stages, each p/v × layers/p deep.

Bubble fraction = (p-1) / (v × m + p - 1) — shrinks by factor v.

Cost: more send/recv operations (v× more boundaries) and more scheduling complexity.

Megatron-LM 2’s 1T-parameter run used interleaved 1F1B with v=4. Real production configs use v=2 to v=4.

Zero-Bubble Pipeline (2023) — the state of the art

Qi, Wan, Huang, Lin (Sea AI Lab), arxiv:2310.01889, published Nov 2023.

Key insight: the backward pass has two independent halves:

  • B (backward-input): compute dX for upstream stage. Blocking — the upstream stage waits for this.

  • W (backward-weight): compute dW. NOT blocking — you can defer it.

Decoupling them lets the scheduler fill bubbles with W-only compute that doesn’t create pipeline dependencies. Result: truly zero pipeline bubbles in synchronous training, without gradient staleness. Modest throughput gain (~15%) over 1F1B, but the memory story is cleaner.

Companion: PipeOffload (2025) leverages the same F/B/W decomposition to offload activations to CPU during idle windows, cutting activation memory further. Modern Megatron includes both.

The scheduling zoo (know the names)

Schedule

Bubble

Memory

Notes

GPipe

(p-1)/(m+p-1)

High (stash all m)

Legacy

1F1B

(p-1)/(m+p-1)

Low (stash p)

Default

Interleaved 1F1B

(p-1)/(vm+p-1)

Low, more comm

Megatron default at scale

Zero-Bubble (ZB-H1/H2)

~0

Same as 1F1B

2023 SOTA

Chimera / bidirectional

~0

2x weights per rank

Niche

The other pain: activation communication

Stage boundary passes activation (b × s × h × bytes) in forward, gradient of same size in backward. For a Llama-70B slice at b=4, s=8192, h=8192, bf16:

4 × 8192 × 8192 × 2 = 512 MB per micro-batch, per boundary, per direction.

With m=32 micro-batches, p=8 stages, that’s ~16 GB sent between each pair per step in each direction. InfiniBand at 400 Gbps = 50 GB/s; so ~0.3s of comm per step, which needs to overlap with compute. This is why compute-comm overlap is the whole game for PP and why PP micro-batch size matters so much.

PP for inference

Uncommon but useful when:

  • Model doesn’t fit in one node and TP=8 is already saturated.

  • Long prefill: chunked prefill can pipeline naturally across stages.

  • Very high throughput needs — pipeline replicas of a big model.

The latency cost for a single decode is severe: p stages of ~2-3 ms each = ~20 ms inter-token latency for p=8. This is why PP for inference is a throughput tool, not a latency tool. TP wins for latency.

vLLM added --pipeline-parallel-size (PP inference) around 2024; SGLang followed. Both are far more mature at TP.

Practical checklist

  1. Balance the stages. Layers are not identical (embedding, lm_head are outliers). Naive equal-partition wastes capacity. Megatron’s --num-layers-per-virtual-pipeline-stage and libraries like nanotron/torchtitan handle this; verify it.

  2. Choose micro-batch count. Rule of thumb: m ≥ 4p. If your global batch is 512 and p=8, use micro-batch size 16 → m=32. Fewer micro-batches = more bubble; more = more memory.

  3. Interleave. v=2 costs little, drops bubble noticeably. v=4 for extreme scale.

  4. Overlap comm. NCCL send/recv should overlap with the next compute chunk. Profile with Nsight Systems and look for send/recv streams idle-time.

  5. Uneven pipelines are fine. Fewer layers on the first/last stages (which carry embedding/lm_head weight) is standard. Auto-partitioners in Megatron do this.

Exercises

  1. Compute bubble fraction for p=4, m=16 GPipe vs 1F1B (same) and interleaved with v=2, m=16.

  2. For Llama-3-405B split across 8 nodes as PP=8, TP=8, compute inter-node send/recv bytes per training step with global batch = 2048, seq = 8192.

  3. Draw the timeline (F = forward micro-batch, B = full backward, W = weight-grad only) for zero-bubble ZB-H1 on p=4, m=8. You should see W blocks filling what 1F1B leaves as gaps.

Reading order

  1. GPipe paper (skim) — the microbatching idea.

  2. PipeDream 1F1B (skim) — the schedule.

  3. Megatron-LM 2 sections 2–4 — interleaved 1F1B, real perf numbers.

  4. Zero-Bubble paper — the modern schedule.

  5. Megatron-LM source (megatron/core/pipeline_parallel/schedules.py) — read the actual scheduling loop. It is not long.