Pipeline parallelism

Deep dive

How it works

In: a global batch split into m micro-batches, a model cut into p stages · Out: one synchronous optimizer step with correct gradients · State: per stage: its layer weights and the activations of in-flight micro-batches · Cost: bubble (p−1)/(m+p−1) of the step; b·s·h activation bytes per micro-batch per stage boundary

Notation: p stages, m micro-batches, forward time tf and backward tb per micro-batch per stage (tb ≈ 2tf).

GPipe (all F, then all B):                     1F1B (warmup, steady 1F1B, cooldown):
s0 F0F1F2F3 . . . . . . B0B1B2B3               s0 F0F1F2F3   B0F4  B1F5 ...
s1   F0F1F2F3 . . . . B0B1B2B3                 s1   F0F1F2 B0F3  B1F4 ...
s2     F0F1F2F3 . . B0B1B2B3                   s2     F0F1 B0F2 B1F3 ...
s3       F0F1F2F3 B0B1B2B3                     s3       F0B0F1B1F2B2 ...

1F1B: stage s runs min(p−s−1, m) warmup forwards, then alternates F and B, then drains the remaining Bs. Its makespan equals GPipe's, (m+p−1)(tf+tb), but at most p−s activations are live.
Interleaved 1F1B (Megatron): each GPU holds v non-contiguous chunks (e.g. layers 1-5 and 41-45). The bubble becomes (p−1)/(v·m). Requires m to be a multiple of p.
Zero bubble (Qi et al.): split B into B_input (needed by the previous stage, on the critical path) and B_weight (deferrable). Schedule the B_weight work into the bubble. ZB-H1 roughly thirds the 1F1B bubble at equal memory; ZB-H2 reaches near zero with more memory.
DualPipe (DeepSeek-V3): feeds micro-batches from both ends, and overlaps forward/backward compute with expert-parallel all-to-all.

Code

A schedule simulator: generates GPipe and 1F1B per-stage op lists, executes them under the data dependencies, and measures makespan, bubble and peak live activations.

#[derive(Clone, Copy)]
enum Op { F(usize), B(usize) }

fn gpipe(_p: usize, m: usize, _s: usize) -> Vec<Op> {
    (0..m).map(Op::F).chain((0..m).map(Op::B)).collect()
}

fn one_f_one_b(p: usize, m: usize, s: usize) -> Vec<Op> {
    let warm = (p - s - 1).min(m);
    let mut ops: Vec<Op> = (0..warm).map(Op::F).collect();
    let (mut f, mut b) = (warm, 0);
    while f < m { ops.push(Op::F(f)); f += 1; ops.push(Op::B(b)); b += 1; }
    while b < m { ops.push(Op::B(b)); b += 1; }
    ops
}

/// Returns (makespan, peak live activations on stage 0).
fn simulate(p: usize, m: usize, tf: f64, tb: f64, sched: fn(usize, usize, usize) -> Vec<Op>) -> (f64, usize) {
    let plans: Vec<Vec<Op>> = (0..p).map(|s| sched(p, m, s)).collect();
    let (mut fdone, mut bdone) = (vec![vec![None::<f64>; m]; p], vec![vec![None::<f64>; m]; p]);
    let (mut pc, mut free, mut live, mut peak) = (vec![0; p], vec![0.0f64; p], vec![0usize; p], vec![0usize; p]);
    let mut remaining: usize = plans.iter().map(Vec::len).sum();
    while remaining > 0 {
        let mut progressed = false;
        for s in 0..p {
            let Some(&op) = plans[s].get(pc[s]) else { continue };
            let ready = match op {                       // data dependencies
                Op::F(i) => if s == 0 { Some(0.0) } else { fdone[s - 1][i] },
                Op::B(i) => if s == p - 1 { fdone[s][i] } else { bdone[s + 1][i] },
            };
            let Some(ready) = ready else { continue };
            let start = ready.max(free[s]);               // stage runs its list in order
            match op {
                Op::F(i) => { free[s] = start + tf; fdone[s][i] = Some(free[s]); live[s] += 1; }
                Op::B(i) => { free[s] = start + tb; bdone[s][i] = Some(free[s]); live[s] -= 1; }
            }
            peak[s] = peak[s].max(live[s]);
            pc[s] += 1; remaining -= 1; progressed = true;
        }
        assert!(progressed, "schedule deadlocks");
    }
    (free.iter().cloned().fold(0.0, f64::max), peak[0])
}

fn main() {
    let (tf, tb) = (1.0, 2.0);
    for (p, m) in [(4, 8), (8, 32)] {
        for (name, sched) in [("gpipe", gpipe as fn(usize, usize, usize) -> Vec<Op>), ("1f1b", one_f_one_b)] {
            let (t, mem) = simulate(p, m, tf, tb, sched);
            let bubble = 1.0 - m as f64 * (tf + tb) / t;
            println!("p={p} m={m} {name:5}: makespan {t:5.1}  bubble {bubble:.3}  stage-0 live acts {mem}");
            assert!((bubble - (p - 1) as f64 / (m + p - 1) as f64).abs() < 1e-9);
        }
    }
}

Output: both schedules have a bubble of 0.273 at p=4, m=8 and 0.179 at p=8, m=32. Stage 0 keeps 8 or 32 live activations under GPipe, but only 4 or 8 under 1F1B.

Complexity and constants

Bubble (p−1)/(m+p−1). Interleaved: about (p−1)/(v·m+p−1).
Activation memory at stage 0: m·A_stage for GPipe, p·A_stage for 1F1B, where A_stage is the activation size of one micro-batch for one stage. Combine with checkpointing on the early stages, which hold the most.
P2p traffic per micro-batch per boundary: b·s·h·2 bytes each way. Example: b = 1, s = 4096, h = 8192 gives 64 MB, about 1.3 ms at 50 GB/s IB. That overlaps with the next micro-batch's compute when send/recv are asynchronous.
Megatron PTD-P (TP 8 within a node, PP across nodes, DP outermost) reached 52% of A100 peak at 3072 GPUs.

Efficiency

Cost: the bubble wastes (p−1)/(m+p−1) of each step (18% at p = 8, m = 32). Each stage boundary moves b·s·h·2 bytes per micro-batch each way: 64 MB for b = 1, s = 4096, h = 8192, about 1.3 ms at 50 GB/s, hidden behind the next micro-batch's compute. Megatron PTD-P reached 52% of A100 peak at 3072 GPUs (Narayanan et al.). Bound: idle time from the bubble and stage imbalance, not bandwidth.
Levers:
- more micro-batches (m ≥ 4p) — bubble falls toward 0 (formula above) — batch large enough to split
- interleaved chunks (v per GPU) — bubble ÷ v, over 10% throughput at equal memory, more p2p traffic (Narayanan et al. 2021) — fast links between stages
- zero-bubble B_input/B_weight split — up to 15% throughput over 1F1B at similar memory, up to 30% with more memory (Qi et al. 2024) — schedule-aware framework
- 1F1B ordering — p live micro-batches of activations instead of m — memory-bound stage 0
- balance by measured stage time — the slowest stage sets the pace — heavy embedding or head stages
Naive → standard → state of the art:
- one micro-batch, stages in turn · only 1 of p GPUs busy
- GPipe or 1F1B with m ≥ 4p · bubble (p−1)/(m+p−1)
- interleaved 1F1B, zero-bubble or DualPipe schedules · bubble near zero, more memory or traffic

Implementation notes

Balance stages by measured time, not layer count. Embedding and LM head (large vocab) stages are heavier; give them fewer transformer blocks.
Send activations asynchronously (isend/irecv, batched p2p) and double-buffer receive tensors.
Gradient all-reduce or reduce-scatter for DP can start per stage as soon as that stage's last backward finishes, overlapping the cooldown.
With interleaving, use m divisible by p and a deterministic chunk-to-GPU mapping. Megatron requires num_microbatches % pipeline_parallel_size == 0.
Tied embeddings (input embedding on stage 0, LM head on stage p−1) need an extra all-reduce of their gradients between the first and last stages.

Pitfalls

Choosing p too large for the batch: at m = p the bubble is about 50%.
Unbalanced stages: one stage 20% slower makes every stage 20% slower.
Deadlocks from mismatched send/recv order between neighbors. Generate both sides from the same schedule function.
BatchNorm or other cross-micro-batch statistics break micro-batch independence.
Loss scaling and averaging must use the total number of micro-batches (divide by m) on the last stage only.

Variants and alternatives

GPipesimplest, synchronous, m activations per stage.

1F1B / PipeDream-Flushsame bubble, bounded memory. The standard choice.

Interleaved 1F1Bsmaller bubble, more communication.

Zero Bubble (ZB-H1/H2), DualPipenear-zero bubble, more complex schedules.

Asynchronous PipeDream (weight stashing)no flush, stale weights. Rarely used for LLMs.

ZeRO-3/FSDPoften simpler up to about 100B if the interconnect allows it.

In production

Megatron-LM/Core (1F1B, interleaved), DeepSpeed pipeline engine, PyTorch torch.distributed.pipelining (GPipe, 1F1B, interleaved, zero-bubble schedules; used by TorchTitan), DeepSeek-V3 (DualPipe), GPipe in Lingvo, PaLM-era JAX pipelines.

Implementations

name	lang	repo	install	verdict	licence	health	as_of	notes
Megatron-LM / Megatron Core	Python / CUDA	NVIDIA/Megatron-LM	`pip install megatron-core`	—	—	—	—	1F1B and interleaved 1F1B
PyTorch pipelining	Python / C++	pytorch/pytorch	`pip install torch`	—	—	—	—	`torch.distributed.pipelining`: GPipe, 1F1B, interleaved, zero-bubble schedules
DeepSpeed	Python / CUDA	deepspeedai/DeepSpeed	`pip install deepspeed`	—	—	—	—	pipeline engine, combines with ZeRO-1
TorchTitan	Python	pytorch/torchtitan	`pip install torchtitan`	—	—	—	—	reference training stack using PyTorch pipelining
DualPipe	Python	deepseek-ai/DualPipe	n/a (clone the repo)	—	—	—	—	bidirectional schedule from DeepSeek-V3

Sources

Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019. paper
Narayanan et al., PipeDream: Generalized Pipeline Parallelism for DNN Training, SOSP 2019. paper
Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021. paper · code
Qi, Wan, Huang, Lin, Zero Bubble Pipeline Parallelism, ICLR 2024. paper · code
DeepSeek-AI, DeepSeek-V3 Technical Report, 2024. paper · code
PyTorch documentation, Pipeline Parallelism (torch.distributed.pipelining). docs