Pipeline Parallelism

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Naive Pipeline (All-Forward Then All-Backward)
5. GPipe: Micro-Batching
6. 1F1B Schedule
7. Interleaved Stages
8. Bubble Fraction Analysis
9. Combining PP with TP and DP
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive the bubble fraction for naive pipeline parallelism.
2. Explain how micro-batching reduces the bubble.
3. Derive the 1F1B schedule and its memory advantage.
4. Compute the optimal number of micro-batches for a given pipeline depth.
5. Design a 3D parallelism strategy combining PP, TP, and DP.

Notation

• $P$ — number of pipeline stages
• $M$ — number of micro-batches
• $t_f, t_b$ — time for one micro-batch forward/backward per stage
• $B_{\text{bubble}}$ — bubble fraction (idle time / total time)

Core Intuition

If a model has 32 layers and you have 4 GPUs, assign 8 layers per GPU (stages). Data flows through stages sequentially: stage 1 → stage 2 → stage 3 → stage 4 (forward), then backward. The problem: while stage 4 computes, stages 1-3 are idle. Micro-batching splits the batch into smaller chunks that flow through the pipeline simultaneously, reducing idle time.

Naive Pipeline (All-Forward Then All-Backward)

Schedule: All micro-batches go forward through all stages, then all go backward.

Timeline for $P=4$, $M=1$:

• Time 1: Stage 1 forward → Stages 2,3,4 idle.
• Time 2: Stage 2 forward → Stages 1,3,4 idle.
• Time 3: Stage 3 forward → Stages 1,2,4 idle.
• Time 4: Stage 4 forward → Stages 1,2,3 idle.
• Time 5-8: Backward (same pattern reversed).

Bubble time: $(P-1)$ idle slots per phase (forward and backward).

Bubble fraction: $(P-1)/(P-1+M) \to (P-1)/P$ for $M=1$. With $P=4$: 75% idle!

GPipe: Micro-Batching

Idea: Split batch into $M$ micro-batches. Pipeline them through stages:

Forward: Micro-batch 1 enters stage 1, then stage 2, ... while micro-batch 2 enters stage 1, etc.

Bubble fraction:

For large $M$: $B_{\text{bubble}} \approx (P-1)/M$. With $P=4, M=32$: bubble = $9.4\%$.

Drawback: Must store activations for all $M$ micro-batches simultaneously (for backward pass). Memory: $O(M \times \text{activations per micro-batch})$.

1F1B Schedule

One Forward, One Backward: Alternate forward and backward micro-batches after the pipeline fills up.

Schedule (after warm-up):

1. Warm-up: stages complete one forward pass each (pipeline filling).
2. Steady state: each stage alternates 1 forward + 1 backward.
3. Cool-down: complete remaining backward passes.

Memory advantage: At any time, each stage holds at most $P$ micro-batch activations (not $M$). Memory: $O(P)$ vs $O(M)$ for GPipe.

Bubble fraction: Same as GPipe: $(P-1)/M$. But much less memory.

Interleaved Stages

Idea (Megatron-LM v3): Assign multiple non-consecutive stages to each GPU.

With $P=4$ GPUs and 16 model chunks:

• GPU 0: chunks 0, 4, 8, 12
• GPU 1: chunks 1, 5, 9, 13
• GPU 2: chunks 2, 6, 10, 14
• GPU 3: chunks 3, 7, 11, 15

Benefit: Each GPU processes $v = 4$ virtual stages. Bubble fraction becomes:

With $v=4$: $4\times$ less bubble for same $M$.

Cost: More communication (activations must transfer between GPUs more frequently).

Bubble Fraction Analysis

Optimal micro-batches: To achieve $<5\%$ bubble with $P$ stages:

For $P=8$: need $M > 140$ micro-batches.

Global batch size constraint: $B_{\text{global}} = M \times B_{\text{micro}} \times \text{DP degree}$. Large $M$ requires large global batch size.

Tradeoff: More micro-batches → less bubble, but larger batch sizes (potentially hurting convergence) and more pipeline flushes.

Combining PP with TP and DP

3D parallelism (Megatron-DeepSpeed):

• TP (tensor): within a node, 4-8 GPUs.
• PP (pipeline): across nodes, 4-16 stages.
• DP (data): across PP replicas, 8-64x.

Total GPUs: $N_{\text{total}} = \text{TP} \times \text{PP} \times \text{DP}$.

Example: Training a 175B model on 512 GPUs:

• TP = 8 (within node, 8 GPUs per node).
• PP = 8 (8 pipeline stages across 8 nodes).
• DP = 8 (8-way data parallel replicas).
• Total: $8 \times 8 \times 8 = 512$ GPUs.

Common Pitfalls

Pitfall 1. Unbalanced stages. If layers have different compute costs (e.g., attention layers with different sequence lengths), stages must be balanced by compute, not just layer count.

Pitfall 2. Using too few micro-batches. With $M < P$: bubble > 50%. Always ensure $M \gg P$.

Pitfall 3. Ignoring activation memory in 1F1B. While 1F1B saves memory vs GPipe, each stage still holds $P$ activations. For very deep pipelines ($P=16$), this is significant.

Summary

• Pipeline parallelism splits layers across GPUs; data flows sequentially.
• Bubble overhead: $(P-1)/M$ — reduced by more micro-batches.
• 1F1B: Same efficiency as GPipe but $O(P)$ memory instead of $O(M)$.
• Interleaved stages: $v\times$ less bubble by assigning multiple chunks per GPU.
• 3D parallelism: TP (intra-node) × PP (inter-node) × DP (replicas).

Exercises

Exercise 1. For $P=8, M=64$: compute the bubble fraction for GPipe, 1F1B, and interleaved ($v=4$).

Exercise 2. Derive the optimal pipeline depth $P$ for 64 GPUs, given TP=8 and the constraint that bubble < 5%.

Exercise 3. Compute the activation memory per GPU for 1F1B vs GPipe with $P=4, M=32, d=4096, T=2048$.

Exercise 4. Design a 3D parallelism configuration for training a 530B model on 1024 A100 GPUs.

Exercise 5. Explain why pipeline parallelism is preferred over tensor parallelism for inter-node communication.