Pipeline parallelism

Overview

Pipeline parallelism is a way to train or serve a very deep neural network by placing different groups of consecutive layers on different GPUs. Each group is called a stage. The batch is cut into small micro-batches, and the micro-batches flow through the stages like items on an assembly line. While one stage works on micro-batch number two, the next stage can work on micro-batch number one, so several devices stay busy at once.

The method lets a model be larger than the memory of any single device, and it needs much less communication than tensor parallelism, because only the activations at stage boundaries are passed along. Its main cost is idle time, known as the pipeline bubble, at the start and end of each step. Good schedules and enough micro-batches keep the bubble small. In practice pipeline parallelism is combined with data parallelism and tensor parallelism. For tabular foundation models it matters when the network is deep or when training must span several servers.

What it is

Pipeline parallelism is a kind of model parallelism that splits the model between layers, which is why it is also called inter-layer parallelism. A model with 48 layers could be split into four stages of 12 layers each. Stage 1 lives on GPU 1, stage 2 on GPU 2, and so on.

Think of a laundry service with four machines in a row: wash, dry, fold, and pack. If the owner waited for one basket to finish all four steps before starting the next, three machines would always sit idle. Instead, as soon as the first basket moves from the washer to the dryer, a second basket enters the washer. After a short start-up period, all four machines work at once. The start-up and the final emptying are the bubble.

The ingredients

• Stages. Contiguous groups of layers, each on its own device or device group.
• Micro-batches. Small slices of the batch that travel through the pipeline one after another.
• Activations. Outputs passed forward from one stage to the next.
• Activation gradients. Gradients passed backward from one stage to the previous one.
• Schedule. The order in which forward and backward passes of micro-batches are run.

Why not just split the layers naively?

If you place layers on different GPUs and run one batch through, only one GPU works at a time. The others wait. This naive model parallelism has terrible utilization. Pipelining fixes it by keeping several micro-batches in flight.

Where it came from

Splitting a network by layers across devices is an old idea in neural network systems, used when a single accelerator could not hold the model. The inefficiency of the naive approach was the obvious problem.

Two papers around 2018 and 2019 made pipelining practical for deep learning. A paper known as GPipe showed how to split a mini-batch into micro-batches, run them through the stages, and apply one synchronous update at the end, with recomputation of activations to save memory. Another line of work, known as PipeDream, interleaved forward and backward passes so that devices were busy more of the time, and explored keeping several versions of the weights.

Later work refined the schedules. The one-forward-one-backward schedule, often called 1F1B, starts backward passes as early as possible so that stored activations can be freed sooner. Interleaved schedules assign several small, non-adjacent chunks of layers to each device, which reduces the bubble further at the price of more communication. Libraries such as Megatron-LM, DeepSpeed, and the pipeline features of PyTorch put these ideas into practice.

How it works

Step 1: Partition the model into stages

The model is cut into stages with the aim of equal work per stage. The slowest stage sets the speed of the whole pipeline, just as the slowest machine in the laundry line does. Partitioning can be done by:

• Counting layers (works when layers are alike, as in a transformer stack).
• Counting parameters.
• Measuring the time of each layer on a profile run and balancing the totals.

Embedding layers at the start and output layers at the end are often heavier or lighter than middle layers, so the first and last stages may need fewer or more transformer blocks to balance.

Step 2: Split the batch into micro-batches

The global batch of B examples is cut into m micro-batches. Each micro-batch has B divided by m examples. The number m is a key setting. Larger m means a smaller bubble, but smaller micro-batches that may use the device less efficiently.

Step 3: Send micro-batches through the forward pass

Micro-batch 1 enters stage 1. When stage 1 finishes it, the activations at the boundary are sent to stage 2, and stage 1 starts micro-batch 2. The activations move forward one stage at a time.

Step 4: Run the backward pass

Once the last stage finishes the forward pass of a micro-batch and computes the loss, the backward pass begins. Gradients flow in reverse, from the last stage to the first. Each stage computes its weight gradients for that micro-batch and passes the gradient with respect to its input to the previous stage.

Step 5: Accumulate gradients and update

Weight gradients from all micro-batches are added together on each stage. After the last backward pass has finished, every stage applies the optimizer step to its own layers. Because the gradients are summed across all micro-batches, the result equals training on the full batch.

The schedules

Fill and drain (GPipe style). All forward passes of all micro-batches run first. Then all backward passes run. It is simple. Its drawback is that activations for every micro-batch must be stored until the backward pass, which uses a lot of memory, so it is paired with recomputation.

One-forward-one-backward (1F1B). After a short warm-up, each stage alternates one forward pass and one backward pass. Each stage holds activations for only a limited number of micro-batches, roughly the number of stages, so memory is lower. The bubble size is the same as in the fill-and-drain schedule for the same m.

Interleaved 1F1B. Each device holds several model chunks, for example layers 1-3 and 13-15 on one device. Micro-batches pass through the device more than once. The bubble shrinks in proportion to the number of chunks, and communication increases in proportion too.

The size of the bubble

With p stages and m micro-batches, the fraction of time lost to the bubble in the basic schedules is about (p - 1) divided by (m + p - 1). If p is 4 and m is 4, the fraction is 3 divided by 7, about 43 percent. If m is 32, it is 3 divided by 35, about 9 percent. This is why the number of micro-batches should be several times larger than the number of stages.

Memory and recomputation

Each stage must keep activations of in-flight micro-batches until their backward pass. Two things reduce this load.

• The 1F1B schedule limits the number of in-flight micro-batches.
• Gradient checkpointing discards most activations and recomputes them during the backward pass.

What stays the same

Pipeline parallelism with a synchronous update gives the same result as single-device training, apart from small numeric differences. It does not change the model, the loss, or the batch semantics. Asynchronous variants change the semantics and are used less often.

Key concepts and terms

A worked example

A timeline with four stages and four micro-batches

Let F1 mean the forward pass of micro-batch 1, and B1 the backward pass. Time moves left to right. Assume each forward or backward unit takes one slot, and the backward costs the same, to keep the picture simple.

Time:     1   2   3   4   5   6   7   8   9  10  11  12  13  14
Stage 1: F1  F2  F3  F4  .   .   .   .   .   .  B4  B3  B2  B1
Stage 2:  .  F1  F2  F3  F4  .   .   .   .  B4  B3  B2  B1  .
Stage 3:  .   .  F1  F2  F3  F4  .   .  B4  B3  B2  B1  .   .
Stage 4:  .   .   .  F1  F2  F3  F4  B4 B3  B2  B1  .   .   .

Dots are idle slots. The total length is 14 slots, while the useful work per stage is 8 slots. Each stage is idle for 6 of 14 slots, which is about 43 percent, matching the formula (p - 1) divided by (m + p - 1) = 3/7 when the forward and backward times are counted together.

Now raise the micro-batches to 16. The total length becomes 2 times (16 + 3) = 38 slots, with 32 useful, so the idle share is 6/38, about 16 percent. The bubble is still there, but it is a smaller share.

A PyTorch-style sketch

The following sketch shows the idea with a simple manual split on two GPUs. Real implementations overlap the stages using schedules, and they use libraries to do so.

import torch
import torch.nn as nn

class TwoStageNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stage1 = nn.Sequential(
            nn.Linear(128, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU()
        ).to("cuda:0")
        self.stage2 = nn.Sequential(
            nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 1)
        ).to("cuda:1")

    def forward(self, x):
        h = self.stage1(x.to("cuda:0"))
        return self.stage2(h.to("cuda:1"))

model = TwoStageNet()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

x = torch.randn(256, 128)
y = torch.randn(256, 1).to("cuda:1")

n_micro = 8
opt.zero_grad(set_to_none=True)
for xb, yb in zip(x.chunk(n_micro), y.chunk(n_micro)):
    loss = loss_fn(model(xb), yb) / n_micro
    loss.backward()   # gradients accumulate across micro-batches
opt.step()

This code splits the batch into micro-batches and accumulates gradients, but it runs them one after another, so the two GPUs do not truly overlap. A pipeline engine would launch micro-batch 2 on stage 1 while stage 2 handles micro-batch 1. The sketch shows the data movement and the accumulation, which are the same in the real engine.

A configuration sketch

pipeline:
  stages: 4
  micro_batches: 16
  schedule: 1f1b
  partition: balanced_by_time
  activation_checkpointing: true
layout:
  gpus: 32
  tensor_parallel: 4
  pipeline_parallel: 4
  data_parallel: 2        # 4 * 4 * 2 = 32

Here the pipeline has four stages and sixteen micro-batches, so the bubble is about 3/19, roughly 16 percent for the basic schedule, before any interleaving.

Reading the bubble formula

The formula for the idle fraction, (p - 1) divided by (m + p - 1), is worth a closer look.

• With one stage (p = 1) the idle fraction is zero. There is no pipeline.
• With one micro-batch (m = 1) the idle fraction is (p - 1) divided by p. Almost all devices are idle, which is the naive model parallel case.
• As m grows large, the fraction tends to zero.

The table below shows a few values for four stages and for eight stages. The numbers come directly from the formula.

Deeper pipelines need many more micro-batches to stay efficient. That is a real cost, because a larger number of micro-batches means a larger global batch or a smaller micro-batch.

The effect of interleaving

If each device holds v chunks of layers, the bubble shrinks by about a factor of v. With four devices and two chunks each, the pipeline behaves like a pipeline with a bubble of about half the size, though with twice as many boundary crossings. The extra messages are small point-to-point sends, so interleaving is attractive when the links are good enough.

Where activations live

Understanding memory is key to a stable pipeline run. Follow the life of one micro-batch.

1. Stage 1 runs the forward pass and stores the activations it will need for the backward pass.
2. It sends the output to stage 2, which does the same.
3. The activations stay stored on each stage until the gradient for that micro-batch comes back.
4. The backward pass uses and then frees them.

Under 1F1B, stage 1 has up to p micro-batches in flight, and the last stage has about one. So the first stage needs the most activation memory, and the last stage the least. A balanced plan may give the first stage fewer layers to even out total memory.

Handling uneven layers

Real networks are not made of identical layers. A few sources of imbalance:

• Embedding layers with large tables are heavy in memory but light in compute.
• Output heads can be heavy in compute if they produce many classes or bins.
• Attention layers with long contexts cost more than feed-forward layers.
• Layers with different widths cost different amounts.

The partitioning goal is to equalize time while respecting memory. A simple dynamic program can choose cut points to minimize the longest stage. Many libraries offer an automatic mode that uses a profile.

Pipeline parallelism at inference

Serving a pipelined model works in a similar way, with requests acting as micro-batches.

• Throughput is high when many requests are in flight, because all stages stay busy.
• Latency for one request is the sum of all stage times plus transfer times, so it is not lower than on one device.
• It suits offline or batch scoring jobs more than low-latency single requests.

For latency-sensitive serving, tensor parallelism is usually the better tool, with pipelining added only if the model still does not fit.

A reasoning example for cross-server training

Suppose a model needs 32 devices to hold its state, the servers have 8 devices each, and the network between servers is much slower than the links inside a server.

1. Use tensor parallelism of degree 8, so each layer spreads over one whole server.
2. Use pipeline parallelism of degree 4, so the model's layers are spread over four servers.
3. Only boundary activations cross the slow network, and only between neighboring servers.
4. If more servers are available, replicate this arrangement with data parallelism and shard the optimizer state.

The same logic applies at any scale: put the chattiest communication on the fastest links.

Debugging a pipeline

When a pipeline run hangs or gives odd results, work through these points.

1. Check that every stage receives the right shapes. A mismatch in boundary shapes is the most common cause of a hang.
2. Check that every micro-batch is processed once. Counting forward and backward passes per stage reveals dropped or duplicated work.
3. Run with two stages and two micro-batches before scaling up.
4. Compare the loss with a non-pipelined version on identical data.
5. Print per-stage timings to find the slow stage.
6. Look for stages that are waiting for a send or receive that never arrives.

Asynchronous pipelines

Some designs avoid the flush by letting a stage start the next batch before the previous one has fully finished. This removes the bubble but means different stages use different versions of the weights for the same micro-batch. Remedies include keeping several versions of weights, which costs memory, or accepting some staleness. Synchronous schedules are favored today because they preserve exact training semantics and are easier to debug.

Strengths and limitations

Strengths

• Lets models far larger than one device be trained by dividing them across devices.
• Communication is light: only boundary activations and their gradients move, and only point to point.
• Works across servers with ordinary network links, unlike tensor parallelism.
• Weight memory drops in proportion to the number of stages.
• Synchronous schedules give results equal to single-device training.
• Combines freely with data, tensor, and sequence parallelism.

Limitations

• The pipeline bubble wastes some time at the start and end of each step.
• Stage imbalance wastes time too, since the slowest stage sets the pace.
• Needs many micro-batches, which may shrink the per-device batch below the efficient size.
• Activation memory can be high for the first stages, which hold in-flight micro-batches longest.
• Layers that share weights across stages, such as tied embeddings, require extra synchronization.
• Code is more complex than data parallelism, and debugging spans multiple devices.
• Models with branching or dynamic control flow are harder to partition.

How it compares with alternatives

Pipeline parallelism versus sharding

Sharded data parallel methods keep every device busy with the full model computation on its own data, but they gather weights for each layer. Pipeline parallelism never gathers weights, because each device owns its layers permanently. It moves activations instead. When the network between servers is slow, moving activations at a few boundaries is often cheaper than gathering the weights of every layer.

Pipeline parallelism versus tensor parallelism

Tensor parallelism has no bubble but needs fast links. Pipeline parallelism has a bubble but tolerates slow links. For this reason the two are placed at different levels of the hardware: tensor parallel inside a server, pipeline parallel across servers.

Combining methods

A standard large-scale layout multiplies three degrees: tensor parallel degree times pipeline parallel degree times data parallel degree equals the number of devices. Mixed precision is applied everywhere. See Megatron-LM and Colossal-AI for libraries that support the combination.

Role in tabular foundation models

Tabular foundation models learn from raw tables: rows and columns, mixed types, missing values, and often no useful column names. They are pre-trained on large numbers of synthetic or real tables and then used with in-context learning, where labeled rows are shown at inference and the model predicts new rows without weight updates. See /model/architecture and /model/research for the broader picture.

When a tabular model needs pipelining

Most tabular models are modest in size, and plain data parallelism is enough. Pipeline parallelism becomes relevant when:

1. The model is deep, with many alternating layers for row and column attention.
2. The total parameter count outgrows the memory of the devices that are available, as can happen with a foundation model trained for broad coverage.
3. Training must span several servers connected by ordinary network links.
4. The same model is used for several heads, such as imputation, anomaly detection, and sector classification, which increases the number of layers and parameters to store.

Natural stage boundaries in a tabular network

A tabular network has a clear structure that suggests where to cut:

• Encoder stage. Per-type embeddings for numeric, categorical, and date columns, plus mask handling for missing values.
• Body stages. Alternating blocks of attention over columns and attention over rows. See axial attention.
• Head stage. Output heads for prediction, imputation, anomaly scores, and semantic sector heads.

The body is made of repeated, similar blocks, so splitting it evenly is easy. The encoder and the heads differ in cost, so the first and last stages may need adjusted loads.

Micro-batches of tables

For tabular pre-training, one sample is typically one table, a set of labeled rows plus query rows. Micro-batches are then groups of tables. Because tables vary in size, micro-batches with very different total cell counts cause jitter in stage times. Grouping tables of similar size into micro-batches keeps the stages in step. This is the same reason sequence models bucket texts by length.

Long contexts for in-context learning

When the context holds many example rows, a single table can be large. Pipeline parallelism helps by spreading the layers across devices, so each device stores fewer weights and has more memory for the activations of the long context. Combined with FlashAttention and gradient checkpointing, this makes longer contexts feasible. At inference, the same split lets a large model serve requests, though the added latency of passing through stages should be weighed against throughput.

Missing values, mixed types, and no column names

The encoder stage handles mixed types and masks, so it is typically the first stage. Since column names are optional or absent, the model relies on column statistics and learned relations. These all sit inside the early stages and require no special treatment in the pipeline. The mask tensor must be passed along, or recomputed, where later stages need it. Any side input that later stages use should be sent with the activations, which slightly increases the boundary message size.

Pre-training objectives and heads

Pre-training uses objectives such as masked cell modeling and imputation, described in imputation networks. The loss is computed in the last stage, so the final stage holds the loss heads and sends gradients backward. If several heads are trained together, the last stage may become heavy. In that case it can be given fewer body layers, or the heads can be spread over the last two stages.

Continual learning

When a model is updated on new tables, each stage stores the penalty terms or replay information for its own layers. For methods like elastic weight consolidation and synaptic intelligence, importance values live with the weights they describe, so they are naturally partitioned by stage. For experience replay, the replay tables are fed to the first stage like any other micro-batch, with no change to the stage structure. Checkpoints should record the stage boundaries so a model can be reloaded under a different layout.

Anomaly detection and sector classification

Scoring heads sit on the final stage. Their inputs are the representations produced by the whole pipeline, so small numeric differences between layouts can shift scores. Test that scores under the pipeline layout match a single-device run on a validation set, especially if thresholds are used. For evaluation and benchmarking, keep the layout fixed. See /model/benchmarks.

When not to use it

For small and mid-sized tabular models, the bubble and the added complexity are not worth it. Start with data parallelism and mixed precision. Add FSDP if memory is short. Move to pipeline parallelism only when the model, even sharded, does not fit or when cross-server links are too slow for sharding.

Practical guidance and best practices

Balance the stages

• Profile each layer before partitioning.
• Aim for equal time per stage, not equal layer count.
• Give the embedding and head stages fewer body layers if they are heavy.
• Re-profile after changing precision or sequence length, since balance can shift.

Choose enough micro-batches

• Use several times more micro-batches than stages.
• Check that the micro-batch is still large enough to use the device well.
• Use gradient accumulation across micro-batches, so the global batch is unchanged.

Pick a schedule

• Use 1F1B to limit activation memory.
• Consider interleaving when the bubble is large and the network can carry the extra messages.
• Keep fill-and-drain for simple setups and for debugging.

Manage memory

• Turn on gradient checkpointing if the first stage runs out of memory.
• Use mixed precision to halve activations and messages.
• Watch the memory on each stage separately, since they are not equal.

Handle shared weights

• Tied input and output embeddings live on different stages and need their gradients synchronized.
• Consider untying them if the synchronization is a nuisance.

Verify equivalence

• Compare the loss over a few steps with a single-device run on the same data.
• Compare gradient norms per layer.
• Re-run the test after changing the stage boundaries.

Plan checkpoints and restarts

• Save each stage's weights along with the layout.
• Provide a tool to merge or re-split stages.
• Save often, since a failure on any stage stops the whole pipeline.

Common mistakes

• Too few micro-batches. The bubble then dominates, and the extra GPUs sit idle.
• Unbalanced stages. One heavy stage makes all others wait.
• Ignoring the first and last stage. Embeddings and heads change the balance.
• Forgetting side inputs. Masks, positions, or padding information needed in later stages must travel with the activations.
• Using pipelining for a model that fits on one device. The bubble costs time with no benefit.
• Running pipeline stages across slow links without measuring. Activations can still be large for long contexts.
• Assuming memory is equal across stages. The first stage holds activations longest and often runs out first.
• Mixing asynchronous schedules without understanding the staleness. Updates computed with old weights change training behavior.
• Not testing equivalence. A partition bug can silently drop or duplicate a layer.
• Placing tensor parallel groups across servers while pipeline stages stay inside a server. This reverses the right order and hurts speed.

Checklist before a long pipeline run

• Stages were balanced by measured time, not by layer count alone.
• The number of micro-batches is several times the number of stages.
• A short equivalence test against one device passed.
• Side inputs such as masks travel with the activations.
• Memory was checked on the first stage, which usually fills first.
• Tied weights are synchronized across the stages that share them.
• Checkpoints record the stage layout.
• Tensor parallel groups sit inside a server, and pipeline stages span servers.

Quick reference of settings

Summary of the trade-off

Pipeline parallelism trades a little idle time for the ability to fit and train a larger model on cheaper links. If the bubble is small and the stages are balanced, the loss in efficiency is modest, and the gain in model size is large.

Frequently asked questions

What is pipeline parallelism in simple terms?

Pipeline parallelism splits a model by layers into stages, places each stage on a different GPU, and sends small pieces of the batch, called micro-batches, through the stages one after another. It works like an assembly line: while one stage handles one piece, the next stage can handle the previous one, so many GPUs work at once.

What is the pipeline bubble?

The pipeline bubble is the idle time at the start and end of each training step, when some stages have no work because the pipeline is still filling or already draining. It shrinks as the number of micro-batches grows relative to the number of stages. Interleaved schedules reduce it further.

How many micro-batches should I use?

Use several times more micro-batches than stages, because the idle fraction is roughly the number of stages minus one divided by the number of micro-batches plus the number of stages minus one. The limit is the efficiency of each micro-batch: if they become too small, the device is underused and the gain disappears.

How is pipeline parallelism different from tensor parallelism?

Pipeline parallelism cuts between layers, so each device owns whole layers and passes activations onward. Tensor parallelism cuts inside layers, so every device takes part in every layer and communicates many times per step. Pipeline parallelism suits slower links and deeper models, while tensor parallelism suits fast links and wide layers.

What is the 1F1B schedule?

One-forward-one-backward is a schedule in which, after a warm-up phase, each stage alternates between a forward pass for one micro-batch and a backward pass for another. This lets stored activations be freed sooner, so memory use is lower than in schedules that run all forward passes first. The bubble size is about the same.

Does pipeline parallelism change training results?

With a synchronous schedule, the gradients from all micro-batches are summed before the update, so the result equals training on the full batch, apart from small numeric differences. Asynchronous schedules change the semantics because updates use stale weights, and they require extra care.

Can pipeline parallelism be combined with data parallelism?

Yes, and large runs almost always do. Each pipeline of stages forms a replica, and several replicas process different data. The gradients for each stage are averaged across replicas. See data parallelism for details of the averaging.

Do tabular foundation models need pipeline parallelism?

Usually not for small or medium models, where one table fits on a device and data parallelism is enough. It becomes useful for very deep or very large models, for runs across several servers, and for long in-context sets of rows that need extra memory for activations. The natural cut points are between the encoder, the attention blocks, and the heads.

How do I choose the number of stages?

Use the smallest number of stages that fits the model in memory, then check the bubble. More stages reduce weight memory per device but add to the bubble and to the number of messages. Balance the stages by measured time, and make sure the number of micro-batches is comfortably larger than the number of stages.

What happens to memory on each stage?

Weight memory is roughly equal across stages if the partition is balanced, but activation memory is not. The first stage holds activations for more in-flight micro-batches than later stages in the 1F1B schedule. Gradient checkpointing and lower precision are the usual remedies. See gradient checkpointing.

Glossary

• 1F1B: A schedule that alternates one forward and one backward pass per stage after warm-up.
• Activation: The output of a layer, passed to the next layer or stage.
• Bubble: Idle time at the start and end of a pipeline step.
• Flush: The end of a step, when all micro-batches have finished and weights update.
• Interleaved schedule: A schedule that gives each device several small chunks of layers.
• Micro-batch: A small slice of the batch that moves through the pipeline.
• Model parallelism: Dividing a model among devices.
• Stage: A group of consecutive layers on one device.
• Stage imbalance: Unequal compute time across stages.
• Tied weights: Parameters shared between two layers, such as input and output embeddings.
• Warm-up: The phase at the start of a step when the pipeline fills.

Summary

Pipeline parallelism divides a model into stages of consecutive layers, places each stage on its own device, and streams micro-batches through the stages so that devices work at the same time. It needs little communication, tolerates ordinary network links, and lets models far larger than one device be trained. The costs are the bubble at the edges of each step, the need for balanced stages, and extra complexity. Enough micro-batches, a good schedule such as 1F1B, and checkpointing keep it efficient.

For tabular foundation models, pipelining is relevant at the large end: deep stacks of row and column attention, big encoders for mixed types and missing values, multiple heads for imputation, anomaly detection, and sector classification, and long in-context sets of rows. Smaller models do well with data parallelism and mixed precision. When pipelining is used, group tables by size into micro-batches, balance the first and last stages, pass side inputs such as masks along, and test equivalence with a single-device run. Related reading: tensor parallelism, FSDP, DeepSpeed ZeRO, Megatron-LM, and /model/architecture.

Related

It belongs to Foundation models from scratch in my stack, under Foundation models and applied ML. Others in the same group: