Pipeline Parallelism — infra.bea.ai

Pipeline parallelism (PP) partitions the model’s layers across GPUs in stages. GPU 0 runs layers 0-9, GPU 1 runs layers 10-19, etc. Data flows through the pipeline like an assembly line — but the “bubble” (idle time) at the start and end is the key challenge.

1. The Core Idea

Stage 0 (GPU 0): Layers 0-9       Stage 1 (GPU 1): Layers 10-19
  Input → [L0-L9] → activations →   [L10-L19] → activations →
                                      ...
Stage 2 (GPU 2): Layers 20-29     Stage 3 (GPU 3): Layers 30-39
  → [L20-L29] → activations →       [L30-L39] → output

Communication: Only point-to-point between adjacent stages — send activations forward, send gradients backward. Much less communication than TP (no AllReduce).

Memory per GPU: Each GPU holds only $L/P$ layers’ weights, gradients, and optimizer states. Memory splits linearly with PP degree.

2. The Bubble Problem

If we process one batch at a time, most GPUs are idle most of the time:

         Micro-batch 0 (forward then backward)
Stage 0: [F0]                         [B0]
Stage 1:      [F0]               [B0]
Stage 2:           [F0]     [B0]
Stage 3:                [F0][B0]
         ^^^^^^^ bubble: GPU 0 idle while other stages run ^^^^^^^

Bubble fraction: With $P$ stages and 1 micro-batch: $(P-1)/P$ . For $P=4$ : 75% of GPU time is idle. Unacceptable.

3. Micro-Batching: Filling the Pipeline

Split the batch into $M$ micro-batches and pipeline them:

GPipe Schedule

Process all micro-batch forwards first, then all backwards:

Stage 0: [F0][F1][F2][F3]                   [B3][B2][B1][B0]
Stage 1:      [F0][F1][F2][F3]         [B3][B2][B1][B0]
Stage 2:           [F0][F1][F2][F3][B3][B2][B1][B0]
Stage 3:                [F0][F1][F2][F3][B3][B2][B1][B0]

Bubble: $(P-1)/(M+P-1)$ . With $M=8, P=4$ : $3/11 = 27\%$ . Better, but still significant.

Memory issue: All $M$ micro-batches’ activations must be stored during forward for backward → $M\times$ activation memory.

1F1B (One Forward, One Backward)

1F1B Schedule

Interleave forward and backward passes:

Stage 0: [F0][F1][F2][F3][B0][F4][B1][F5][B2][F6][B3][B4][B5][B6]
Stage 1:      [F0][F1][F2][B0][F3][B1][F4][B2][F5][B3][B4][B5]
Stage 2:           [F0][F1][B0][F2][B1][F3][B2][F4][B3][B4]
Stage 3:                [F0][B0][F1][B1][F2][B2][F3][B3]

Advantage: Each stage keeps at most $P$ micro-batches’ activations in memory (not $M$ ), because backward frees activations as it goes.

Bubble: Same as GPipe — $(P-1)/(M+P-1)$ . But memory is much better.

Interleaved 1F1B (Virtual Stages)

Each GPU handles multiple non-contiguous stages (e.g., GPU 0 handles stages 0 and 4):

Physical GPU 0: Stage 0 (Layers 0-4) + Stage 4 (Layers 20-24)
Physical GPU 1: Stage 1 (Layers 5-9) + Stage 5 (Layers 25-29)
...

Bubble: $(P-1)/(M \times V + P-1)$ where $V$ = virtual stages per GPU. With $V=2$ : bubble halves.

Tradeoff: More point-to-point communication (stages are non-contiguous), but significantly less bubble time.

4. Communication

Communication	Volume	Frequency
Activations (forward)	$B_{\mu} \times S \times d \times 2$ bytes per micro-batch	Once per stage boundary
Gradients (backward)	Same size	Once per stage boundary

For $B_{\mu}=4, S=2048, d=4096$ , BF16: $4 \times 2048 \times 4096 \times 2 = 64$ MB per transfer.

On InfiniBand (50 GB/s): 64 MB / 50 GB/s = 1.3 ms per transfer. With 4 stages: 3 boundaries × 1.3 ms = 3.9 ms per micro-batch direction. Acceptable on InfiniBand — this is why PP can be inter-node.

5. Pipeline Parallelism in Practice

# Megatron-LM / DeepSpeed
args.pipeline_model_parallel_size = 4    # 4 stages
args.num_micro_batches = 8                # 8 micro-batches per step
args.virtual_pipeline_model_parallel_size = 2  # interleaved

PP uses only point-to-point communication (low bandwidth), unlike TP which needs AllReduce (high bandwidth). This is why PP can go inter-node while TP must stay intra-node.

When to Use PP

Condition	Use PP?
Model fits on one node with TP	No — TP + DP is enough
Model too large for one node	Yes — PP across nodes, TP within node
Communication bandwidth is low	Yes — PP needs only point-to-point (less than TP’s AllReduce)
Need to minimize latency	No — PP adds pipeline startup latency

Typical composition: DP × PP × TP. Example: 256 GPUs = DP=8 × PP=4 × TP=8.

Self-Check

Draw the pipeline bubble for 4 stages with 1 micro-batch
Calculate the bubble fraction for GPipe with P stages and M micro-batches
Explain the memory advantage of 1F1B over GPipe
Explain interleaved stages and how they reduce the bubble
Calculate point-to-point communication volume between PP stages
Explain why PP can be inter-node while TP must be intra-node