Pipeline parallelism — turning a deep Transformer into an assembly line

Yesterday we looked at tensor parallelism: multiple GPUs cooperate on the matrix multiplications inside the same Transformer layer.

Today we split the model in a different direction.

Instead of dividing each layer across GPUs, we put different groups of layers on different GPUs:

This is pipeline parallelism.

The intuition is an assembly line. If a model is too large for one GPU, one device can own the beginning of the network, another the middle, and another the end. But merely splitting the model does not make training fast. To keep all those GPUs busy, we also have to split the batch into smaller microbatches and send them through the stages in a staggered pipeline.

That creates the central tradeoff of pipeline parallelism:

Start with a model too large for one GPU¶

Suppose a Transformer has 48 layers.

A simple two-GPU partition might be:

GPU 0: layers  1–24
GPU 1: layers 25–48

For an input \(x\), GPU 0 computes:

It sends the intermediate activation \(h_{24}\) to GPU 1, which computes:

So the forward pass becomes:

tokens
  │
  ▼
┌──────────────────┐
│ GPU 0            │
│ layers 1–24      │
└──────────────────┘
        │
        │ activations
        ▼
┌──────────────────┐
│ GPU 1            │
│ layers 25–48     │
└──────────────────┘
        │
        ▼
      logits

Unlike tensor parallelism, the two GPUs do not need to cooperate on every matrix multiplication. Each stage can run several complete Transformer layers locally before communicating.

That can substantially reduce the frequency of communication.

But the naive pipeline wastes GPUs¶

Suppose stage 0 takes one unit of time and stage 1 also takes one unit.

If we send an entire batch through at once, the schedule looks like:

time →      1        2

GPU 0:    batch     idle
GPU 1:    idle      batch

GPU 1 cannot begin until GPU 0 produces its activations.

Then GPU 0 sits idle while GPU 1 finishes.

We doubled the number of GPUs but did not double throughput.

The problem is not memory anymore.

It is scheduling.

Microbatching fills the pipeline¶

Suppose the training batch contains 64 sequences.

Instead of processing all 64 as one unit, divide it into four microbatches:

Call them:

M1  M2  M3  M4

GPU 0 processes M1 and sends its activations to GPU 1.

While GPU 1 works on M1, GPU 0 starts M2:

time →       1       2       3       4       5

GPU 0:      M1      M2      M3      M4      idle
GPU 1:      idle    M1      M2      M3      M4

Now both GPUs are busy during most of the schedule.

This is the key idea:

The model is partitioned across space — different GPUs — while the microbatches are staggered across time.

The empty slots are the pipeline bubble¶

Look again at the schedule:

time →       1       2       3       4       5

GPU 0:      M1      M2      M3      M4       -
GPU 1:       -      M1      M2      M3      M4

Those - slots are idle compute.

They form the pipeline bubble.

For a simplified forward-only pipeline with:

• \(p\) pipeline stages,
• \(m\) microbatches,
• equally fast stages,

the pipeline takes approximately:

time slots.

There are \(m\) useful microbatch operations per stage, so a rough utilization is:

For two stages and four microbatches:

If we instead use 32 microbatches:

The startup and shutdown cost still exists, but it is amortized over much more work.

So:

Why not use thousands of tiny microbatches?¶

Because the utilization equation hides another kind of efficiency.

GPUs are exceptionally good at large matrix operations. If a microbatch becomes too small, matrix multiplications may no longer contain enough parallel work to saturate the accelerator.

There is also overhead per microbatch:

kernel launches
activation transfers
scheduler bookkeeping
communication latency

So we have competing effects:

more microbatches
      │
      ├── smaller pipeline bubble
      │
      └── smaller GPU operations and more overhead

The best microbatch size depends on the model, sequence length, accelerator, network, and the other forms of parallelism being used.

This is a recurring theme in distributed LLM systems: maximizing one efficiency metric can make another worse.

Training is harder because backward flows in reverse¶

Everything above described the forward pass.

Training also requires backpropagation.

Forward communication goes:

GPU 0 ── activations ──→ GPU 1

Backward communication goes the opposite way:

GPU 0 ←─ activation gradients ─ GPU 1

For a particular microbatch, GPU 0 cannot finish its backward computation until GPU 1 has run the later layers and propagated the gradient back across the stage boundary.

So a real training pipeline must schedule both forward and backward work while respecting dependencies.

One widely used family of schedules is called 1F1B: after the pipeline warms up, a stage alternates roughly between one forward microbatch and one backward microbatch.

Conceptually:

forward M3
backward M1
forward M4
backward M2
...

The exact timeline depends on the implementation and number of stages, but the goal is simple:

Pipeline parallelism is therefore partly a neural-network partitioning problem and partly a classic scheduling problem.

The slowest stage sets the pace¶

Now suppose a four-stage pipeline takes:

per microbatch.

The first three stages can each produce work every 10 ms.

But stage 4 consumes it only every 25 ms.

Eventually work piles up in front of stage 4 and the faster stages have to wait.

In steady state, throughput is constrained approximately by:

where \(t_i\) is the execution time of stage \(i\).

This means dividing a 48-layer Transformer as:

12 layers | 12 layers | 12 layers | 12 layers

is not necessarily optimal.

Layers can have different costs. The first or last stage may also own embeddings, the output projection, or loss computation.

A better partition could be:

14 layers | 13 layers | 12 layers | 9 layers + output head

if that makes the actual execution times more equal.

The ideal pipeline is not the one with equal numbers of layers.

It is the one with approximately equal stage times.

What actually crosses a pipeline boundary?¶

Suppose the activation tensor between two stages has shape:

Take:

With BF16 activations at 2 bytes per element, one boundary tensor occupies:

That is about:

So a pipeline handoff can move hundreds of megabytes.

Training also sends gradients back across that boundary.

Pipeline communication is therefore not cheap. But compare its pattern with tensor parallelism.

Tensor parallelism may require collective communication repeatedly inside every Transformer layer.

Pipeline parallelism usually communicates only at stage boundaries, after several local layers have run.

A useful approximation is:

tensor parallelism:
frequent, tightly synchronized communication

pipeline parallelism:
less frequent, larger point-to-point transfers

This distinction becomes important when we map parallelism strategies onto real GPU interconnects.

Pipeline parallelism and tensor parallelism compose¶

These techniques divide the model along different axes.

Tensor parallelism splits a layer:

Layer 1: [ GPU 0 | GPU 1 ]
Layer 2: [ GPU 0 | GPU 1 ]
Layer 3: [ GPU 0 | GPU 1 ]

Pipeline parallelism splits depth:

GPU 0: [Layer 1][Layer 2][Layer 3]
                         │
                         ▼
GPU 1: [Layer 4][Layer 5][Layer 6]

Large models often use both.

Suppose we have 64 GPUs and choose:

Then:

One model replica uses 16 GPUs:

pipeline stage 0       pipeline stage 1       stage 2       stage 3

┌──────────────┐       ┌──────────────┐
│ 4-way tensor │ ───→  │ 4-way tensor │ ───→   ...   ───→
│   parallel   │       │   parallel   │
└──────────────┘       └──────────────┘

Then four such model replicas participate in data parallelism.

This combination is often called 3D parallelism:

ZeRO can further shard training state across the data-parallel dimension.

The important point is that these techniques are not competitors. They solve different pieces of the distributed-training problem.

A concrete mental model: a restaurant kitchen¶

Imagine a busy restaurant with four stations:

prep → grill → plating → service

If there is only one order, most stations wait.

The grill cannot work until prep finishes. Plating cannot work until the grill finishes.

But with many orders:

Order 1: prep → grill → plating → service
Order 2:        prep → grill → plating → service
Order 3:               prep → grill → plating → service

all four stations can eventually work at once.

The orders are the microbatches.

The stations are the pipeline stages.

The empty time while the kitchen fills and drains is the pipeline bubble.

And if grilling takes three times longer than every other station, the entire restaurant's throughput becomes limited by the grill.

That is exactly why stage balance matters.

The deeper lesson¶

It is tempting to think that distributed training is mainly about making a model fit:

The model is too large for one GPU, so spread it over several.

But once we do that, a second problem appears.

We must ensure that the expensive hardware we added is actually busy.

Pipeline parallelism makes this especially visible because utilization becomes a timeline:

useful work | useful work | idle | useful work | ...

The central optimization is no longer merely:

It is also:

That gives us today's main principle:

And this sets up the next question naturally. Once tensor parallelism, pipeline parallelism, and data parallelism are all operating together, the physical network connecting the GPUs becomes part of the model architecture in practice: a parallelization strategy that looks perfect on paper can be terrible if its communication pattern does not match the hardware topology.