Pipeline Parallelism: Pumping Data Upstream

In our previous exploration of Tensor Parallelism (TP), we sliced the model horizontally, splitting individual weight matrices across GPUs. While effective, TP performs frequent collectives, so fast links are often valuable; TP can also span nodes when the interconnect and workload justify it.

To scale beyond a single node—to clusters of hundreds or thousands of GPUs—we need a technique that is less "chatty" and more tolerant of the lower bandwidth available over Ethernet or InfiniBand.

Enter Pipeline Parallelism (PP). Instead of slicing the layers themselves, PP slices the model vertically, partitioning the layers of the neural network across multiple devices. It turns your GPU cluster into an assembly line, passing activations down the chain.

This article dissects the hardware mechanics of PP, the scheduling algorithms required to keep the silicon busy, and the critical trade-offs involved.

1. The Hardware Mechanics: Vertical Slicing

TP divides operations within a layer; PP assigns different layers to different stages. Both distribute computation and parameter memory.

The Setup

Imagine a hypothetical 32-layer transformer distributed across four GPUs (stages).

• GPU 0: Layers 0–7
• GPU 1: Layers 8–15
• GPU 2: Layers 16–23
• GPU 3: Layers 24–31

The Communication Primitive: Point-to-Point (P2P)

PP commonly uses point-to-point Send/Recv between stages. TP commonly uses collectives such as All-Reduce; All-Reduce combines tensors and is not a broadcast.

• Forward Pass: GPU 0 computes the activations for Layer 7 and Sends them to GPU 1. GPU 1 Receives them and begins computing Layer 8.
• Backward Pass: GPU 1 computes gradients for Layer 8 and Sends the gradient-with-respect-to-input back to GPU 0.

Hardware Implication: At a simple stage boundary, each microbatch transfers an activation tensor forward and an activation-gradient tensor backward. Their sizes depend on microbatch, sequence length, hidden width and dtype. Relative to TP, the communication cost depends on the workload and topology; PP can be useful across nodes.

2. The "Bubble" Problem: Silicon Idle Time

The fundamental flaw of a naive pipeline is sequential dependency. With one microbatch and a simple sequential schedule, downstream stages wait for the forward activations; backward work proceeds in reverse. The time spent waiting is called the Pipeline Bubble.

To fight this, we inject Micro-Batches. We split the global batch into $M$ small chunks and pump them into the pipeline back-to-back.

Schedule A: GPipe (All-Forward-All-Backward)

This is the simplest implementation.

1. Fill: Pump all $M$ micro-batches through the forward pass.
2. Drain: Once the last micro-batch reaches the end, start the backward passes in reverse order.

The Hardware Cost: Memory. Each stage retains forward state for its microbatches until their backward passes, unless it recomputes activations. Activation memory can therefore grow with microbatch count.

Schedule B: 1F1B (One-Forward-One-Backward)

This schedule is implemented in frameworks such as PyTorch and Megatron-LM. Instead of waiting for all forward passes to finish, a GPU essentially swaps between forward and backward tasks once the pipeline is full (the "steady state").

1. Warmup: GPUs process enough forward passes to fill the pipeline.
2. Steady State: 1 Forward (create new activation) $\rightarrow$ 1 Backward (consume old activation, free memory).
3. Cooldown: Process remaining backward passes.

The Hardware Benefit: By interleaving operations, we drastically reduce the peak memory footprint. The forward state for microbatch $k$ can be freed after its backward pass, rather than keeping all microbatches until the end of the training step.

3. Bare Metal Implementation: The 1F1B Logic

Implementing a pipeline requires managing a state machine on each GPU. The following pseudocode shows the ordering on an intermediate stage. It omits buffer shapes, matching peer schedules, backward bookkeeping, loss handling and synchronization. For executable code, use a supported pipeline runtime.

# Schematic synchronous 1F1B schedule for one intermediate stage.
# All stages must use a matching send/receive order; this is not runnable code.
for each warmup microbatch:
    activation = recv(previous_stage)
    output = local_forward(activation)
    send(output, next_stage)
    retain_backward_state(microbatch)
for each steady-state slot:
    forward_one_new_microbatch()          # receive, compute, send, retain
    gradient = recv(next_stage)           # oldest outstanding microbatch
    input_gradient = local_backward(gradient)
    send(input_gradient, previous_stage)
for each remaining outstanding microbatch:
    backward_one_microbatch()             # receive, compute, send
optimizer_step()                          # after all microbatches, all stages

The schedule is illustrative. In particular, a stage cannot simply issue blocking sends and receives without coordination with neighboring stages.

4. Pros and Cons: When to use PP?

Pros

• Inter-Node Scaling: Because P2P communication is lightweight (only activations at the boundaries), PP can be placed across nodes when the activation transfers and pipeline bubbles fit the network and workload.
• Memory Efficiency: By splitting the model layers, parameter, gradient and optimizer states are distributed by stage. Per-device memory depends on layer sizes, embeddings, tied weights, activations and partition balance; four stages do not automatically use one quarter of total model memory each.

Cons

• The Bubble (Idle Compute): The biggest cost. During warmup and cooldown, GPUs sit idle. If your number of micro-batches is small, the efficiency drops drastically. For a simple balanced pipeline with equal forward and backward stage times, no communication cost and $M$ microbatches on $N$ stages, the idealized utilization is $M/(M+N-1)$. Imbalance and communication lower real utilization.
• Stale Weights / Complexity: Some asynchronous schedules use multiple weight versions. A synchronous, flush-based 1F1B schedule performs the optimizer update after all microbatches of a training step; schedule details matter for equivalence to non-pipelined training.
• Static Shapes: Variable shapes require communicating shape metadata or using fixed-size buffers/padding; runtime support varies.

Summary: The 3D Parallelism Stack

Pipeline Parallelism sits in the middle of the 3D scaling hierarchy:

1. Tensor Parallelism: Inside the node (intra-layer). Fast links are helpful.
2. Pipeline Parallelism: Between nodes (inter-layer). Can cross nodes depending on workload and network.
3. Data Parallelism: Replicates the entire TP+PP cluster to scale batch size.

By combining these, we can train models with trillions of parameters on thousands of GPUs.

Sources

• PyTorch distributed pipeline parallelism
• GPipe paper
• Megatron-LM pipeline scheduling paper

Related

• Tensor Parallelism Deep Dive — the intra-layer half of the stack.
• Expert Parallelism: The Router's Dilemma — how MoE models shard the feed-forward instead.