Pipeline parallelism — splitting Transformer layers across GPUs

Over the last three lessons, we have split LLM training in three different ways.

Data parallelism split the batch across GPUs.

FSDP split the persistent training state across GPUs.

Tensor parallelism split the computation inside a layer across GPUs.

Today we will split along a fourth axis:

This is pipeline parallelism.

The intuition is simple. If a 96-layer Transformer cannot fit comfortably on one GPU, put layers 1–24 on GPU 0, layers 25–48 on GPU 1, and so on.

But that immediately creates a problem: if we process one batch at a time, most GPUs spend most of their time waiting.

The central trick of pipeline parallelism is therefore not merely dividing the layers.

It is dividing the batch into microbatches so different GPUs can work on different microbatches simultaneously.

Start with a four-stage model¶

Imagine a 96-layer Transformer divided across four GPUs:

GPU 0 → layers 1–24
GPU 1 → layers 25–48
GPU 2 → layers 49–72
GPU 3 → layers 73–96

Call these pipeline stages:

For an input activation \(x\), the model computes:

The important point is that GPU 1 cannot begin until GPU 0 produces \(h_1\).

GPU 2 waits for GPU 1.

GPU 3 waits for GPU 2.

So with one batch, execution looks like:

time →

GPU 0: [ stage 0 ]
GPU 1:             [ stage 1 ]
GPU 2:                         [ stage 2 ]
GPU 3:                                     [ stage 3 ]

This saves memory, but it is terrible for utilization.

At any moment, only one GPU is doing useful work.

The assembly-line mental model¶

Imagine building cars.

One station installs the engine.

Another installs the doors.

Another paints the car.

Another performs final inspection.

If you build exactly one car, the stations work sequentially:

car 1:

engine → doors → paint → inspect

Most stations are idle most of the time.

But a factory does not wait for car 1 to finish before starting car 2.

Once car 1 leaves the engine station, car 2 enters it.

Soon:

engine station  → car 4
door station    → car 3
paint station   → car 2
inspection      → car 1

All four stations are busy simultaneously.

Pipeline parallelism does the same thing with microbatches.

Split the batch into microbatches¶

Suppose the training batch is divided into four microbatches:

Instead of sending all of them through stage 0 before doing anything else, we pipeline them.

A simplified forward schedule looks like:

time →       t1    t2    t3    t4    t5    t6    t7

GPU 0       m1    m2    m3    m4
GPU 1             m1    m2    m3    m4
GPU 2                   m1    m2    m3    m4
GPU 3                         m1    m2    m3    m4

At time \(t_4\):

GPU 0 → microbatch 4
GPU 1 → microbatch 3
GPU 2 → microbatch 2
GPU 3 → microbatch 1

All four GPUs are computing.

This is the pipeline's steady state.

But there is still a bubble¶

Look at the beginning:

t1:
GPU 0 busy
GPU 1 idle
GPU 2 idle
GPU 3 idle

Then:

t2:
GPU 0 busy
GPU 1 busy
GPU 2 idle
GPU 3 idle

The pipeline takes time to fill.

At the end, it also takes time to drain.

Those idle regions are called the pipeline bubble.

For \(p\) pipeline stages and \(m\) microbatches, a simple forward-only pipeline needs approximately:

pipeline time slots.

Useful stage-work occupies:

slots per stage.

So an idealized utilization is:

This equation gives us an important intuition.

A concrete example¶

Suppose:

pipeline stages and only:

microbatch.

Then:

That matches our original sequential picture.

Now use four microbatches:

Then:

With 32 microbatches:

So:

The fixed fill-and-drain cost gets amortized across more work.

Training complicates the pipeline¶

So far we have discussed only forward passes.

Training also needs backward passes.

For one microbatch:

forward:

GPU 0 → GPU 1 → GPU 2 → GPU 3

backward:

GPU 3 → GPU 2 → GPU 1 → GPU 0

Backward flows in the opposite direction because each stage needs gradients from the stage after it.

A naive schedule could do all forwards first:

F1 F2 F3 F4

and then all backwards:

B4 B3 B2 B1

But that can require storing many activations while waiting for backward.

Remember from our activation-checkpointing lesson: backward needs information from forward.

If many microbatches are simultaneously "in flight," activation memory can become large.

The 1F1B idea¶

A widely useful scheduling idea is one-forward-one-backward, usually abbreviated 1F1B.

After the pipeline warms up, a stage alternates approximately between:

one forward microbatch
one backward microbatch
one forward microbatch
one backward microbatch
...

Conceptually:

time →

stage 0: F1 F2 F3 F4 B1 F5 B2 F6 B3 ...
stage 1:    F1 F2 F3 B1 F4 B2 F5 B3 ...
stage 2:       F1 F2 B1 F3 B2 F4 B3 ...
stage 3:          F1 B1 F2 B2 F3 B3 ...

Real schedules depend on the implementation, but the important idea is that forward and backward work are interleaved.

This reduces the number of microbatches whose activations must remain live simultaneously.

So pipeline scheduling is not just about utilization.

It is also about memory.

What crosses a pipeline boundary?¶

Suppose GPU 0 owns layers 1–24 and GPU 1 owns layers 25–48.

GPU 1 does not need GPU 0's weights.

It needs the activation produced by layer 24.

If the activation tensor is:

then GPU 0 sends \(H\) to GPU 1.

During backward, GPU 1 sends back the gradient:

So the communication pattern is roughly:

forward:
stage i → activations → stage i+1

backward:
stage i ← activation gradients ← stage i+1

This is quite different from tensor parallelism.

Tensor parallelism may communicate multiple times inside each layer.

Pipeline parallelism communicates mainly at stage boundaries.

Why pipeline parallelism can tolerate slower links better than tensor parallelism¶

Suppose each pipeline stage contains 24 Transformer layers.

A GPU may perform substantial computation before it needs to send one activation tensor to the next stage.

The ratio is roughly:

If stages contain enough compute, communication can be relatively infrequent.

Tensor parallelism is different.

Its GPUs may need collectives every Transformer block.

That gives us a useful topology rule:

fastest links
→ tensor parallelism

somewhat slower links
→ pipeline boundaries can sometimes tolerate them better

This is not absolute, but it explains why large clusters carefully map different parallelism dimensions onto different network levels.

Stage balance matters enormously¶

Suppose four stages take:

stage 0 → 10 ms
stage 1 → 10 ms
stage 2 → 10 ms
stage 3 → 25 ms

The pipeline cannot advance faster than its slowest stage.

In steady state, new work emerges roughly every:

not every 10 ms.

The faster stages repeatedly wait.

This is analogous to the slowest station in a factory assembly line.

If painting takes twice as long as every other operation, cars pile up before painting.

So a good pipeline partition tries to equalize:

not necessarily:

Different layers may have different costs.

Embedding layers, attention variants, MoE layers, and output heads can make equal layer counts quite unequal in execution time.

Memory balance matters too¶

Suppose the compute is perfectly balanced:

stage 0 → 20 ms
stage 1 → 20 ms
stage 2 → 20 ms
stage 3 → 20 ms

but stage 3 contains a giant output vocabulary projection.

Its parameter memory may be much larger.

So pipeline partitioning often considers both:

and:

The best cut points may therefore look irregular:

GPU 0 → layers 1–27
GPU 1 → layers 28–53
GPU 2 → layers 54–78
GPU 3 → layers 79–96 + output head

rather than equal 24-layer chunks.

Pipeline parallelism changes activation lifetime¶

Without pipeline parallelism, a GPU executes all layers for a microbatch.

With pipeline parallelism, a stage only owns activations for its own layers.

That can reduce activation memory per GPU.

But microbatching introduces another dimension.

If a stage has several forward microbatches waiting for their backward passes, it may need to retain activation state for each one:

This is why 1F1B scheduling and activation checkpointing work well together.

Checkpointing reduces activation storage within a microbatch.

Pipeline scheduling controls how many microbatches' activations are live at once.

These techniques attack different factors in the same memory equation.

Pipeline parallelism versus gradient accumulation¶

The word microbatch appeared earlier in gradient accumulation.

The concepts are related but not identical.

Gradient accumulation says:

Process several microbatches before performing one optimizer update.

Pipeline parallelism says:

Keep several microbatches in flight so different pipeline stages remain busy.

Conveniently, the same microbatch decomposition can often serve both purposes.

Suppose:

sequences and:

microbatches contribute to one optimizer update.

Those 16 microbatches can also be the units flowing through the pipeline.

So one decomposition helps both:

gradient accumulation
→ statistical batch size

pipeline parallelism
→ hardware utilization

Combine pipeline and tensor parallelism¶

Now imagine a model with 96 layers running on 32 GPUs.

We might use:

pipeline stages.

Each stage gets eight GPUs.

Inside each stage:

So:

pipeline stage 0
→ layers 1–24
→ tensor-parallel across GPUs 0–7

pipeline stage 1
→ layers 25–48
→ tensor-parallel across GPUs 8–15

pipeline stage 2
→ layers 49–72
→ tensor-parallel across GPUs 16–23

pipeline stage 3
→ layers 73–96
→ tensor-parallel across GPUs 24–31

Within a stage, eight GPUs cooperate on each layer.

Between stages, activations move from one eight-GPU group to the next.

Total GPUs:

Now add data parallelism.

If we have eight copies of that 32-GPU arrangement processing different data:

Then total GPU count becomes:

So:

GPUs.

This simple multiplication is one of the most useful equations for understanding large training layouts.

Add sequence parallelism and sharding, and the picture becomes multidimensional¶

A real large-model training job may combine:

data parallelism
×
FSDP / ZeRO sharding
×
tensor parallelism
×
pipeline parallelism
×
sequence parallelism
×
gradient accumulation
×
activation checkpointing

These techniques are not interchangeable.

Each attacks a different scaling limit.

A useful summary is:

data parallelism
→ divide examples

tensor parallelism
→ divide operations inside a layer

pipeline parallelism
→ divide layers

FSDP
→ divide persistent model state

gradient accumulation
→ divide a batch across time

activation checkpointing
→ divide activation storage across time

Once you see these as orthogonal axes, giant training configurations become much less mysterious.

Pipeline parallelism during inference¶

Pipeline parallelism can also be used for inference.

Suppose different layer ranges live on different GPUs:

token
 ↓
GPU 0: layers 1–24
 ↓
GPU 1: layers 25–48
 ↓
GPU 2: layers 49–72
 ↓
GPU 3: layers 73–96
 ↓
logits

For a single request, latency includes traversing every stage sequentially.

So pipeline parallelism does not magically make one request four times faster.

But with many requests or sequences, different stages can work on different requests simultaneously.

This can improve throughput.

The distinction is important:

That makes workload shape important when choosing an inference architecture.

The deeper mental model: parallelism is choosing an axis to cut¶

Imagine the full training computation as a giant multidimensional object.

It has axes corresponding roughly to:

training examples
sequence positions
hidden dimensions
layers
parameters
time

Different techniques cut along different axes.

Data parallelism cuts:

Tensor parallelism cuts:

Pipeline parallelism cuts:

Sequence parallelism cuts:

FSDP cuts:

The art is choosing cuts whose communication boundaries are cheap enough.

Every cut creates a seam.

Across that seam, something must move:

data parallelism
→ gradients

tensor parallelism
→ partial tensor results

pipeline parallelism
→ activations and activation gradients

FSDP
→ parameters and gradient shards

So distributed training can be understood through one recurring question:

If you know the answer, you can often predict the performance bottleneck.

The progression from one GPU to a cluster¶

We can now take a Transformer that conceptually looks like:

tokens
 ↓
layer 1
 ↓
layer 2
 ↓
...
 ↓
layer 96
 ↓
logits

and map it onto hundreds or thousands of GPUs.

First, split layers into pipeline stages.

Then split each large layer across tensor-parallel GPUs.

Then replicate those groups across different training examples.

Then shard training state where replication wastes memory.

Then use microbatches to keep the pipeline full.

The mathematical model remains one Transformer.

The physical machine may look like:

                    DATA PARALLEL COPIES

        replica 0          replica 1          replica 2
           │                  │                  │

PP0   [TP TP TP TP]      [TP TP TP TP]      [TP TP TP TP]
           ↓                  ↓                  ↓
PP1   [TP TP TP TP]      [TP TP TP TP]      [TP TP TP TP]
           ↓                  ↓                  ↓
PP2   [TP TP TP TP]      [TP TP TP TP]      [TP TP TP TP]
           ↓                  ↓                  ↓
PP3   [TP TP TP TP]      [TP TP TP TP]      [TP TP TP TP]

What looks like one model in PyTorch has become a distributed computation graph spread across a data center.

That gives us today's main principle:

But the assembly line only works well when enough microbatches keep it full and each stage takes roughly the same amount of time.

So the broader lesson from the last few days is becoming clear:

Scaling an LLM is not simply about adding GPUs. It is about partitioning the computation so that thousands of GPUs can remain busy without spending their time waiting for one another.