Pipeline Parallelism › Pipeline Schedules

The order in which forward and backward passes are scheduled across stages and micro-batches is called the pipeline schedule. Different schedules make different tradeoffs between bubble fraction, activation memory, and implementation complexity. The two most important schedules in practice are GPipe and 1F1B, with the interleaved 1F1B schedule providing a further refinement.

The schedule determines how busy each device is, how much memory it requires, and how long the pipeline takes to produce a gradient update. These tradeoffs are non-trivial: a schedule that minimizes the bubble fraction might require enormous activation memory, making it impractical for large models. A schedule that minimizes memory might introduce communication complexity that reduces efficiency. Understanding the specific tradeoffs of each schedule is essential for selecting the right one for a given training configuration, cluster topology, and model architecture.

GPipe (Google's Pipeline for Machine Learning) was one of the earliest systematic frameworks for pipeline parallelism in deep learning. The paper, published by researchers at Google Brain in 2019, demonstrated that models with billions of parameters could be trained reliably using pipeline parallelism with gradient checkpointing. Its schedule is conceptually straightforward: all $M$ micro-batches complete their forward passes before any backward passes begin. This "flush" design is easy to reason about, simple to implement, and was used to train some of the earliest large language models.

The full GPipe schedule for $D$ stages and $M$ micro-batches proceeds as follows. In the forward phase, micro-batches $B_1$ through $B_M$ are injected sequentially into the pipeline. Stage 1 processes $B_1$, then $B_2$, then $B_3$, and so on, passing activations downstream as each micro-batch completes. By the time stage $D$ has processed $B_M$, all activations for all micro-batches are sitting in GPU memory across all stages. Only then does the backward phase begin. Stage $D$ computes gradients for $B_M$, then $B_{M-1}$, and so on in reverse order. Gradients flow backward through the pipeline until all micro-batches have completed their backward passes and all parameter gradients have been accumulated.

The bubble fraction for GPipe is:

$$\text{bubble fraction}_{\text{GPipe}} = \frac{D - 1}{M + D - 1}$$

where:

• $D$: number of pipeline stages
• $M$: number of micro-batches

Note that this formula uses $D-1$ in the numerator rather than $2(D-1)$. This is because in the GPipe flush-based schedule, the startup bubble and the drain bubble are shared: the same pipeline states that are idle during the forward-phase ramp-up become idle again during the backward-phase drain, but the two phases are not counted separately in the same way as a fully interleaved schedule. The effective bubble fraction formula is $(D-1)/(M+D-1)$, which is equivalent to the general formula when you account for the specific timing of the flush.

For $D = 4$ and $M = 8$, this gives $(4-1)/(8+4-1) = 3/11 \approx 0.27$. For $D = 4$ and $M = 32$, it gives $3/35 \approx 0.086$. As $M$ grows large, the bubble fraction approaches zero.

The critical weakness of GPipe is its peak activation memory. Because all $M$ micro-batch activations across all $D$ stages must be stored simultaneously (they are all needed for the eventual backward phase), the total activation memory scales as:

$$\text{activation memory}_{\text{GPipe}} = M \times D \times A_{\text{layer}}$$

where:

• $M$: number of micro-batches
• $D$: number of pipeline stages
• $A_{\text{layer}}$: activation memory per layer per micro-batch

For large $M$ and $D$, this can easily exceed available GPU memory. Consider a GPT-3-scale model with $D = 8$ stages and $M = 32$ micro-batches. Each activation tensor has shape $(B_{\mu}, T, H) = (4, 2048, 12288)$ in bfloat16, which is about 201 MB. The total activation memory is $32 \times 8 \times 201 \text{ MB} \approx 51 \text{ GB}$, which on top of the parameter and optimizer state memory would exceed any single GPU.

GPipe addresses this through gradient checkpointing (also called activation recomputation). Instead of storing all intermediate activations during the forward pass, only certain checkpoint activations are kept. When the backward pass needs an intermediate activation that was not stored, the forward pass is rerun from the nearest checkpoint to recompute it. This reduces activation memory dramatically at the cost of additional computation.

In GPipe, gradient checkpointing is applied at the stage granularity: each stage stores only its input activation (the tensor it received from the previous stage) and recomputes all intermediate activations within the stage during the backward pass. With this approach, the activation memory per stage is just $A_{\text{stage-input}}$ regardless of how many layers are in the stage or how many micro-batches are in flight. The recomputation cost is roughly one additional forward pass for every backward pass, which adds about 30-33% to total compute time (since the backward pass typically costs about twice as much as the forward pass, and you are adding one more forward pass).

The simplicity of GPipe made it the natural first implementation choice, and it demonstrated that pipeline parallelism was practical at scale. The activation memory limitations, combined with the compute overhead of gradient checkpointing, motivated the development of more memory-efficient schedules.

The 1F1B (One Forward, One Backward) schedule, introduced by the Megatron-LM project at NVIDIA, addresses the memory problem of GPipe by interleaving forward and backward passes rather than separating them into two distinct phases. The core idea is that a device should begin its backward pass as soon as it has accumulated enough in-flight forward passes to make the backward pass possible, rather than waiting for all $M$ micro-batches to complete their forward phases first.

In the steady-state of a 1F1B schedule, each device alternates between one forward-pass micro-batch and one backward-pass micro-batch. At any point in time, each device is always working on one computation, and the peak number of activations stored in memory is bounded by the number of stages $D$ rather than the number of micro-batches $M$.

The 1F1B schedule proceeds in three distinct phases.

The warmup phase fills the pipeline: stage 1 processes micro-batches $B_1$ through $B_D$ in the forward direction, and stage $D$ performs its first backward pass after receiving $B_D$ from all prior stages. This warmup phase is where the pipeline bubble accumulates, since early stages are busy on their forward passes while late stages are still waiting for activations to arrive. The warmup phase lasts for $D - 1$ time steps of startup latency.

The steady-state phase then alternates forward and backward passes, maintaining full pipeline utilization. Every stage is either processing a forward pass or a backward pass at every time step. In the steady state, stage $k$ forwards micro-batch $B_{k+D}$ while simultaneously backward-passing micro-batch $B_k$. The stage is never idle during the steady state, as there is always either a new forward micro-batch arriving or a backward gradient ready to process.

The drain phase at the end completes the remaining backward passes as the pipeline empties. After the last forward pass ($B_M$ reaching stage $D$), no new forward passes enter, and the pipeline drains as gradients flow backward through each stage. The drain phase mirrors the startup bubble in duration: it also occupies $D - 1$ time steps.

The bubble fraction for 1F1B is identical to GPipe:

$$\text{bubble fraction}_{\text{1F1B}} = \frac{D - 1}{M + D - 1}$$

The advantage of 1F1B is not a smaller bubble but a much smaller peak activation memory. Because each stage begins processing backward passes before accumulating many pending forward-pass activations, the maximum number of in-flight activations at any stage is bounded by $D$:

$$\text{activation memory}_{\text{1F1B}} = D \times A_{\text{layer}}$$

where:

• $D$: number of pipeline stages
• $A_{\text{layer}}$: activation memory per layer per micro-batch

Comparing to GPipe's $M \times D \times A_{\text{layer}}$, the 1F1B schedule saves a factor of $M$ in activation memory. For a typical training run with $M = 32$ and $D = 8$, 1F1B uses 32 times less activation memory than naive GPipe without checkpointing. This allows 1F1B to handle significantly larger $M$ without running out of memory, which in turn enables smaller micro-batches (for better tensor core utilization) while still achieving a low bubble fraction.

The reason for the memory difference becomes clear when you examine the schedule visually. In GPipe, by the time the backward phase begins, every stage has already processed all $M$ micro-batches in the forward direction and is holding all of their activations in memory simultaneously. In 1F1B, a stage begins its backward pass much sooner. Stage $D$ begins its first backward pass after receiving only its $D$-th micro-batch's activations, not after receiving all $M$ micro-batches. As each backward pass completes at stage $D$, the corresponding forward activations can be freed. The memory at any stage is therefore bounded by the maximum number of micro-batches simultaneously "in flight" at that stage, which is $D$, not $M$.

This is the key algorithmic insight of 1F1B: by scheduling backward passes early, you free activation memory early, keeping the peak memory footprint proportional to pipeline depth rather than to the number of micro-batches. For practical training runs where $M$ can be tens or hundreds of micro-batches, this memory savings is what makes 1F1B the default schedule for large model training.

The interleaved schedule is a further refinement of 1F1B introduced by Megatron-LM to reduce the bubble fraction at the cost of increased communication. Standard 1F1B assigns a contiguous block of $L/D$ layers to each device. The interleaved schedule instead assigns $V$ smaller, non-contiguous blocks of $L/(D \times V)$ layers each to every device. A single device might hold layers 1-4 and layers 33-36, with another device holding layers 5-8 and 37-40, and so on.

This arrangement means that each micro-batch traverses the pipeline $V$ times instead of once, because it must visit each device multiple times, once per model chunk assigned to that device. Each traversal is called a pipeline pass. The bubble fraction for the interleaved schedule improves by a factor of $V$ compared to standard 1F1B:

$$\text{bubble fraction}_{\text{interleaved}} = \frac{1}{V} \cdot \frac{D - 1}{M + D - 1}$$

where:

• $V$: number of model chunks per device (also called virtual pipeline stages)
• $D$: number of devices
• $M$: number of micro-batches

With $V = 2$, the interleaved schedule halves the bubble fraction. With $V = 4$, it reduces the bubble to a quarter. The intuition is that by splitting the model into more, smaller chunks, the pipeline fills up faster relative to the total amount of work. The startup ramp covers the same number of time steps, but those time steps represent a smaller fraction of the total work because each "pass" through the pipeline covers only $1/V$ of the model's layers. The effective depth of each pipeline pass is shorter, so the bubble overhead is smaller relative to the useful computation.

The tradeoff is that each device now sends and receives activation tensors $V$ times per micro-batch instead of once, so inter-device communication volume increases by a factor of $V$. For $V = 2$, the device sends activations at four different points in the forward pass (instead of two: once when passing from chunk 1 to the next device's chunk 2, once again when passing from chunk 3 to the next device's chunk 4, and similarly in the backward pass). Whether this tradeoff is favorable depends critically on the ratio of communication bandwidth to compute throughput. On systems with fast interconnects, such as NVLink between GPUs within the same node, the extra communication overhead may be small relative to the bubble reduction gains. On systems with slower interconnects, such as InfiniBand between nodes, the added communication can become a bottleneck that negates the bubble benefits.

The interleaved schedule also imposes a constraint that the number of micro-batches $M$ must be at least $D \times V$ for full pipeline efficiency, and clean scheduling requires $M$ to be a multiple of $D$. In large-scale training runs where $M$ is already large, these constraints are not binding. In smaller experiments or when memory constraints limit $M$, the constraint can be inconvenient.

The three schedules form a clear progression in terms of capabilities and implementation complexity. Standard GPipe establishes the baseline: full forward phase followed by full backward phase, simple to implement, large activation memory. Standard 1F1B matches GPipe's bubble fraction but dramatically reduces activation memory by interleaving. Interleaved 1F1B goes further, reducing the bubble fraction at the cost of more communication.

Choosing between them in practice comes down to the constraints of the specific hardware setup. For a cluster with fast NVLink interconnects where communication cost is low, interleaved 1F1B with $V = 2$ or $V = 4$ is often the best choice, as the bubble reduction improves hardware utilization without a significant communication penalty. For clusters with slower interconnects where communication is a bottleneck, standard 1F1B may be preferable. GPipe is primarily used today when simplicity is the priority, such as for initial prototyping or for smaller models where activation memory is not a binding constraint.