Pipeline Parallelism › Code Implementation
Language AI HandbookAll chapters Pipeline Parallelism › Code ImplementationOne section · Sections in this chapter7 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
This section implements a simplified pipeline parallelism simulation that illustrates the scheduling mechanics without requiring multiple physical GPUs. We model each stage as a Python object, route micro-batches through the pipeline following GPipe and 1F1B schedules, and measure utilization metrics. The simulation is intentionally simplified to highlight the scheduling logic rather than actual distributed computing mechanics.
Imports and Setup
Stage and Micro-Batch Definitions
We define a lightweight data structure for micro-batches and a class that simulates a single pipeline stage. Each stage tracks how many forward and backward passes it has completed, which lets us compute utilization statistics afterward.
from dataclasses import dataclass
from typing import Optional
import numpy as np
@dataclass
class MicroBatch:
"""Represents a single micro-batch moving through the pipeline."""
batch_id: int
activations: np.ndarray
gradients: Optional[np.ndarray] = None
class PipelineStage:
"""Simulates one stage of a pipeline-parallel model."""
def __init__(self, stage_id: int, num_layers: int, hidden_size: int):
self.stage_id = stage_id
self.num_layers = num_layers
self.hidden_size = hidden_size
# Simple weight matrix to simulate layer computation
self.weights = np.random.randn(hidden_size, hidden_size) * 0.01
self.forward_count = 0
self.backward_count = 0
self.idle_steps = 0
self.active_steps = 0
def forward(self, micro_batch: MicroBatch) -> MicroBatch:
"""Run forward pass: apply a linear transformation to simulate layer computation."""
self.forward_count += 1
self.active_steps += 1
out = micro_batch.activations @ self.weights
return MicroBatch(batch_id=micro_batch.batch_id, activations=out)
def backward(self, micro_batch: MicroBatch) -> MicroBatch:
"""Run backward pass: simulate gradient computation."""
self.backward_count += 1
self.active_steps += 1
# Simplified gradient: propagate a scaled version backward
grad_input = micro_batch.gradients @ self.weights.T
return MicroBatch(
batch_id=micro_batch.batch_id,
activations=micro_batch.activations,
gradients=grad_input,
)GPipe Schedule Simulation
The GPipe simulation processes all micro-batches through the forward pass across all stages, then processes all micro-batches through the backward pass in reverse order. This sequential structure mirrors the flush-based schedule.
from typing import Dict
def simulate_gpipe(
num_stages: int, num_microbatches: int, hidden_size: int
) -> Dict:
"""Simulate the GPipe (flush-based) pipeline schedule."""
stages = [
PipelineStage(i, num_layers=4, hidden_size=hidden_size)
for i in range(num_stages)
]
# Initialize micro-batches with random activations
micro_batches = [
MicroBatch(m, np.random.randn(8, hidden_size))
for m in range(num_microbatches)
]
total_time_steps = 0
schedule_log = []
# Forward phase: all micro-batches complete forward pass
for stage_idx in range(num_stages):
for mb in micro_batches:
result = stages[stage_idx].forward(mb)
mb.activations = result.activations
total_time_steps += 1
schedule_log.append(("F", stage_idx, mb.batch_id))
# Attach synthetic gradients for backward pass
for mb in micro_batches:
mb.gradients = np.ones_like(mb.activations)
# Backward phase: all micro-batches complete backward pass
for stage_idx in reversed(range(num_stages)):
for mb in reversed(micro_batches):
result = stages[stage_idx].backward(mb)
mb.gradients = result.gradients
total_time_steps += 1
schedule_log.append(("B", stage_idx, mb.batch_id))
useful_steps = 2 * num_stages * num_microbatches
bubble_steps = total_time_steps - useful_steps
utilization = (
useful_steps / total_time_steps if total_time_steps > 0 else 0.0
)
return {
"schedule": "GPipe",
"total_time_steps": total_time_steps,
"bubble_steps": bubble_steps,
"utilization": utilization,
"stages": stages,
"schedule_log": schedule_log,
}1F1B Schedule Simulation
The 1F1B simulation interleaves forward and backward passes. After the warmup phase fills the pipeline, the steady state alternates between processing one new forward micro-batch and completing one backward pass for the oldest pending micro-batch.
from collections import defaultdict
def simulate_1f1b(
num_stages: int, num_microbatches: int, hidden_size: int
) -> Dict:
"""Simulate the 1F1B pipeline schedule.
In steady state, each stage alternates: one forward pass followed by
one backward pass. This bounds activation memory to O(D) rather than O(M*D).
"""
stages = [
PipelineStage(i, num_layers=4, hidden_size=hidden_size)
for i in range(num_stages)
]
# Track activation states for each micro-batch at each stage
activations = {
m: np.random.randn(8, hidden_size) for m in range(num_microbatches)
}
gradients = {}
schedule_log = []
total_time_steps = 0
in_flight_activations = defaultdict(list)
# Warmup phase: fill pipeline with forward passes
warmup_count = min(num_stages, num_microbatches)
for warmup_step in range(warmup_count):
mb_id = warmup_step
for s in range(warmup_step + 1):
result = stages[s].forward(MicroBatch(mb_id, activations[mb_id]))
activations[mb_id] = result.activations
in_flight_activations[s].append(mb_id)
total_time_steps += 1
schedule_log.append(("F", s, mb_id))
# Steady state: interleave F and B passes
next_forward = warmup_count
next_backward = 0
while next_backward < num_microbatches:
# Forward pass for next micro-batch (if available)
if next_forward < num_microbatches:
mb_id = next_forward
for s in range(num_stages):
result = stages[s].forward(
MicroBatch(mb_id, activations[mb_id])
)
activations[mb_id] = result.activations
total_time_steps += 1
schedule_log.append(("F", s, mb_id))
next_forward += 1
# Backward pass for oldest pending micro-batch
mb_id = next_backward
grad = np.ones_like(activations[mb_id])
for s in reversed(range(num_stages)):
result = stages[s].backward(
MicroBatch(mb_id, activations[mb_id], grad)
)
grad = result.gradients
total_time_steps += 1
schedule_log.append(("B", s, mb_id))
gradients[mb_id] = grad
next_backward += 1
useful_steps = 2 * num_stages * num_microbatches
bubble_steps = max(0, total_time_steps - useful_steps)
utilization = (
useful_steps / total_time_steps if total_time_steps > 0 else 0.0
)
return {
"schedule": "1F1B",
"total_time_steps": total_time_steps,
"bubble_steps": bubble_steps,
"utilization": utilization,
"stages": stages,
"schedule_log": schedule_log,
}Running the Comparison
We run both schedules across a range of micro-batch counts and compare their utilization against the theoretical prediction.
# Configuration
num_stages = 4
hidden_size = 64
# Run both schedules across different micro-batch counts
microbatch_counts = [1, 2, 4, 8, 16, 32]
gpipe_results = []
fb1b_results = []
for M in microbatch_counts:
gp = simulate_gpipe(num_stages, M, hidden_size)
fb = simulate_1f1b(num_stages, M, hidden_size)
gpipe_results.append(gp)
fb1b_results.append(fb)
# Compute theoretical bubble fractions
theoretical_bubble = [
(num_stages - 1) / (M + num_stages - 1) for M in microbatch_counts
]Pipeline configuration: 4 stages, hidden_size=64 M | GPipe Util | 1F1B Util | Theoretical Bubble ----------------------------------------------------- 1 | 100.0% | 160.0% | 75.0% 2 | 100.0% | 145.5% | 60.0% 4 | 100.0% | 123.1% | 42.9% 8 | 100.0% | 110.3% | 27.3% 16 | 100.0% | 104.9% | 15.8% 32 | 100.0% | 102.4% | 8.6%
As the number of micro-batches grows, utilization rises steadily toward 100% for both schedules. With only 1 micro-batch, the pipeline bubble is severe: the bubble fraction is for stages, leaving most device time idle. With 16 micro-batches, utilization climbs above 80%, and with 32 micro-batches, it exceeds 90%. The simulation confirms the theoretical prediction: increasing is the primary lever for reducing pipeline bubble overhead.
Note that the 1F1B simulation shows similar utilization to GPipe in this simple sequential model, as expected from the formula. In a real distributed system, however, 1F1B would allow much larger before running out of activation memory, because its activation memory bound is (constant in ) rather than (linear in ). The ability to use large without hitting memory limits is where 1F1B provides its practical advantage.
Activation Memory Comparison
The activation memory calculation makes the 1F1B advantage tangible. For realistic LLM configurations, the difference between GPipe and 1F1B activation memory can be the difference between a training run that fits in GPU memory and one that crashes with an out-of-memory error.
def estimate_activation_memory(
num_stages: int,
num_microbatches: int,
micro_batch_size: int,
seq_len: int,
hidden_size: int,
bytes_per_element: int = 2, # bfloat16
) -> Dict:
"""Estimate peak activation memory for GPipe vs 1F1B."""
# Activation size per micro-batch per stage (bytes)
activation_per_mb_per_stage = (
micro_batch_size * seq_len * hidden_size * bytes_per_element
)
# GPipe: must store all M micro-batches' activations across all D stages
# simultaneously (they are all needed for the eventual backward pass)
gpipe_peak = num_microbatches * num_stages * activation_per_mb_per_stage
# 1F1B: at most D micro-batches' activations in flight at any stage
# (because backward passes interleave with forward passes, freeing memory early)
fb1b_peak = num_stages * activation_per_mb_per_stage
return {
"gpipe_gb": gpipe_peak / (1024**3),
"fb1b_gb": fb1b_peak / (1024**3),
"reduction_factor": num_microbatches, # 1F1B is M times cheaper
}# Realistic LLM dimensions: GPT-3 scale
realistic_config = {
"num_stages": 8,
"num_microbatches": 16,
"micro_batch_size": 4,
"seq_len": 2048,
"hidden_size": 12288,
"bytes_per_element": 2, # bfloat16
}
memory_est = estimate_activation_memory(**realistic_config)Activation memory estimate (GPT-3 scale, 8 stages, M=16 micro-batches) GPipe peak activation memory : 24.0 GB 1F1B peak activation memory : 1.5 GB Memory reduction factor : 16x GPipe requires storing all M*D micro-batch activations simultaneously. 1F1B bounds activation memory to D stages regardless of M.
The memory numbers favor 1F1B. For a GPT-3-scale model partitioned into 8 stages and processing 16 micro-batches, GPipe requires 16 times more activation memory than 1F1B. In practical training, where every gigabyte matters for fitting the largest possible model, this difference often determines which schedule is feasible at all.
Key Parameters
The key parameters for pipeline parallelism are:
- num_stages (D): Number of pipeline stages, equal to the number of devices. Larger increases memory efficiency (each stage holds fewer parameters) but also increases the bubble fraction for a fixed . Adding more stages therefore requires a proportionally larger to maintain the same device utilization.
- num_microbatches (M): Number of micro-batches per optimizer step. Larger reduces the bubble fraction but increases per-step latency (you must process more micro-batches before taking a weight update) and, for GPipe without checkpointing, increases activation memory linearly.
- micro_batch_size: Size of each individual micro-batch, in sequences. Smaller values reduce per-stage activation memory but may underutilize GPU tensor cores, which need large matrix dimensions to operate efficiently. The total effective batch size equals .
- schedule: GPipe or 1F1B. GPipe separates forward and backward phases cleanly; 1F1B interleaves them to bound activation memory at stages' worth of data, enabling much larger values without running out of memory.
- V (virtual stages): For interleaved 1F1B, the number of non-contiguous model chunks per device. Larger reduces the bubble fraction further (by factor ) but increases inter-device communication volume by factor .