Learn AI Create the Future — with Pranay
Main site

Multi-GPU and Distributed Training

Pipeline and tensor parallelism

When one layer is too big for one GPU, you stop splitting the batch and start splitting the model — lengthwise into pipeline stages, or crosswise through the middle of each matrix multiply.

Read these first

On this page 5
  1. Why both exist
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Two ways to cut a model that will not fit: across the assembly line, or down the middle of one machine.

Picture a car factory. The car is too big for one workshop.

Pipeline parallelism is the assembly line. Workshop 1 fits the engine, passes the car to workshop 2 for the doors, then workshop 3 paints it. Each workshop owns different steps.

Tensor parallelism is two people painting one car at the same time, one on the left side and one on the right. They share a single step, splitting the work inside it.

Why both exist

The assembly line has an ugly moment at the start of the day. Workshop 3 has nothing to paint until workshops 1 and 2 have finished the first car. Everyone downstream stands idle. That idle time is called the bubble.

The cure is to keep many cars moving at once. Feed small cars in quickly and every workshop stays busy — but you can never remove the bubble at the start and the end.

Side-by-side painting has no bubble. Instead the two painters must constantly check they are matching each other's work. That checking is a conversation, and conversations across machines are slow. So the side-by-side approach is used inside one machine, where the cables between GPUs are fast.

How it works

pipeline (split the layers):
   gpu0: layers 1-3  ->  gpu1: layers 4-6  ->  gpu2: layers 7-9
         data flows left to right; each GPU owns different layers

tensor (split inside a layer):
   gpu0: left half of layer 4's numbers   \
                                            }- combine after every layer
   gpu1: right half of layer 4's numbers  /

A real example you have seen

A restaurant kitchen. The tandoor, the curry station and the plating counter form a pipeline; the first order takes a while, then plates come out steadily. Two cooks chopping opposite halves of the same sack of onions are tensor parallel.

Remember this

  • Pipeline = different GPUs own different layers; the cost is idle time called the bubble.
  • Tensor = different GPUs own pieces of the same layer; the cost is constant chatter.
  • Both exist for when the model is too big, not to go faster on a model that fits.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Two of the three demonstrations below are genuinely runnable without special hardware. The third needs several GPUs and is presented as code only. Distributed pieces were captured on CPU with the gloo backend and two processes.

The simplest model parallelism: put halves on different devices

split_model.py
import torch
import torch.nn as nn

# two "devices": on a real multi-GPU box these are cuda:0 and cuda:1
dev_a = "cpu"
dev_b = "cuda:0" if torch.cuda.is_available() else "cpu"

class SplitNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.first_half = nn.Sequential(nn.Linear(16, 32), nn.ReLU()).to(dev_a)
        self.second_half = nn.Sequential(nn.Linear(32, 4)).to(dev_b)

    def forward(self, x):
        h = self.first_half(x.to(dev_a))
        return self.second_half(h.to(dev_b))       # the hand-off between devices

model = SplitNet()
out = model(torch.randn(2, 16))
out.sum().backward()                               # autograd follows across devices

print("halves live on:", dev_a, "and", dev_b)
print("output device:", out.device, "shape:", tuple(out.shape))
print("gradient reached the first half?",
      model.first_half[0].weight.grad is not None)
Output
halves live on: cpu and cuda:0
output device: cuda:0 shape: (2, 4)
gradient reached the first half? True

That output is environment-dependent: it was captured on a machine with one CUDA GPU, so dev_b resolved to cuda:0. On a CPU-only laptop both halves land on cpu and the device column reads cpu — the script still runs and the last line still prints True.

The point survives either way. Autograd crosses devices on its own. A .to(device) call is a differentiable operation, so the backward pass follows the tensor home without any help from you.

This is model parallelism at its most basic, and it is also the slowest possible version. While second_half computes, dev_a does nothing at all.

Fixing the idle time: the pipeline bubble, measured

pipeline_bubble.py
def timeline(stages, microbatches):
    """Forward-only GPipe schedule: stage s starts microbatch m at time s + m."""
    length = stages + microbatches - 1
    rows = []
    for s in range(stages):
        row = ["."] * length
        for m in range(microbatches):
            row[s + m] = str(m)
        rows.append("".join(row))
    busy = stages * microbatches
    return rows, busy / (stages * length)

for micro in (1, 4):
    rows, use = timeline(stages=4, microbatches=micro)
    print(f"4 GPUs, {micro} microbatch(es) -- each column is one time slot")
    for i, r in enumerate(rows):
        print(f"  gpu{i}: {r}")
    print(f"  GPUs busy {use:.0%} of the time, idle {1 - use:.0%}\n")
Output
4 GPUs, 1 microbatch(es) -- each column is one time slot
  gpu0: 0...
  gpu1: .0..
  gpu2: ..0.
  gpu3: ...0
  GPUs busy 25% of the time, idle 75%

4 GPUs, 4 microbatch(es) -- each column is one time slot
  gpu0: 0123...
  gpu1: .0123..
  gpu2: ..0123.
  gpu3: ...0123
  GPUs busy 57% of the time, idle 43%

The dots are the bubble. With one batch, four GPUs are 75% idle — you bought three cards to achieve nothing. Chopping the batch into four microbatches lifts usage to 57%, and the diagonal shape shows why: the ramp at each end never disappears.

Raise microbatches to 16 and usage reaches 84%. The rule of thumb is to keep microbatches at four or more times the stage count, and the limit is memory, because every in-flight microbatch holds its activations.

Tensor parallelism, done by hand

tensor_parallel.py
import torch
import torch.nn as nn
import torch.distributed as dist

dist.init_process_group(backend="gloo")
rank, world = dist.get_rank(), dist.get_world_size()

torch.manual_seed(0)
full = nn.Linear(8, 6, bias=False)      # the layer we are splitting
x = torch.randn(2, 8)                   # every rank holds the same input
reference = full(x)

# column parallel: rank r keeps only its slice of the OUTPUT features
cols = 6 // world
mine = full.weight[rank * cols:(rank + 1) * cols, :].clone()
partial = x @ mine.T                    # shape (2, cols) -- a piece of the answer
print(f"rank {rank}: my weight slice {tuple(mine.shape)}, my output {tuple(partial.shape)}")

pieces = [torch.zeros_like(partial) for _ in range(world)]
dist.all_gather(pieces, partial)        # stitch the pieces back into one answer
combined = torch.cat(pieces, dim=1)

if rank == 0:
    print("unsplit layer output :", [round(v, 4) for v in reference[0].tolist()])
    print("stitched from shards :", [round(v, 4) for v in combined[0].tolist()])
    print("same answer?", torch.allclose(reference, combined, atol=1e-6))

dist.destroy_process_group()
bash
torchrun --nproc_per_node=2 tensor_parallel.py
Output
rank 0: my weight slice (3, 8), my output (2, 3)
rank 1: my weight slice (3, 8), my output (2, 3)
unsplit layer output : [-0.3064, 0.0628, -0.4779, 0.2413, -0.3609, 0.0032]
stitched from shards : [-0.3064, 0.0628, -0.4779, 0.2413, -0.3609, 0.0032]
same answer? True

That is tensor parallelism, in eight lines and with no library. Each rank stored half the weight matrix, did half the multiply, and one all_gather reassembled an answer identical to the unsplit layer.

Notice what it costs: a collective call per layer, per forward pass. Ten layers means ten collectives. That is why tensor parallelism stays inside one machine, on NVLink, and pipeline parallelism spans machines.

Splitting output features is column parallel. Splitting input features is row parallel, and it ends with all_reduce rather than all_gather. Transformers alternate the two — column for the first feed-forward matrix, row for the second — so a whole block costs one collective instead of two.

The library versions

Real jobs use PyTorch's own implementations rather than hand-written slicing. Both need several GPUs, so no output is shown here:

library_apis.py — needs a real multi-GPU setup
from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel
from torch.distributed.pipelining import pipeline, ScheduleGPipe

# tensor parallel: name which layers are split which way
parallelize_module(block, tp_mesh, {
    "attn.qkv": ColwiseParallel(),
    "attn.out": RowwiseParallel(),
})

# pipeline parallel: split into stages, then pick a schedule
pipe = pipeline(model, mb_args=(example_microbatch,), split_spec={"layer4": "beginning"})
stage = pipe.build_stage(stage_index=rank, device=device)
schedule = ScheduleGPipe(stage, n_microbatches=8)

Common mistakes

Reaching for model parallelism first. It is the last resort. Try mixed precision, then gradient checkpointing, then FSDP. Only when a single layer cannot fit does splitting the layer become necessary.

One microbatch in a pipeline. The measured 75% idle figure above is what that costs.

Tensor parallelism across machines. A collective per layer over Ethernet turns a training run into a network benchmark. Keep the tensor-parallel group inside one node.

Unbalanced pipeline stages. The slowest stage sets the pace for every other. Split by measured time per layer, not by layer count — embedding and output layers are usually much heavier than the blocks between them.

Forgetting that the batch is still split too. Real large-scale training composes data, tensor and pipeline parallelism at once. Your world size is the product of the three, not the sum.

Try it yourself

Run pipeline_bubble.py with stages=8 and microbatches of 1, 8 and 32. Plot the three usage numbers. Then work out how many microbatches you need for 90% usage with 8 stages, and check your answer with the script.

What to learn next

Researcher — Mathematics and papers.

The bubble, exactly

For a GPipe-style schedule with $S$ stages and $M$ microbatches, and assuming equal stage times, the pipeline occupies $S + M - 1$ time slots to do $S \times M$ slots of work. The bubble fraction is

$$ \text{bubble} = \frac{S - 1}{S + M - 1} $$

where $S$ is the stage count and $M$ the microbatch count — matching the 75% and 43% figures printed above for $S=4$, $M \in {1, 4}$. The fraction falls as $O(S/M)$, so utilisation is bought with microbatches, and microbatches are paid for in stored activations: GPipe holds up to $M$ microbatches' activations at the first stage.

1F1B (one-forward-one-backward, from PipeDream) has the identical bubble fraction but caps in-flight activations at $S$ rather than $M$, which is why it is the default in every serious implementation. Interleaved 1F1B (Narayanan et al., 2021) assigns $v$ non-contiguous chunks per device, cutting the bubble to $\frac{S-1}{v(S+M-1)}$ at the cost of $v$ times the point-to-point traffic. Zero-bubble schedules split the backward pass into input-gradient and weight-gradient halves, and can drive the bubble to near zero when stage times are well balanced.

Tensor parallelism in a transformer

Megatron-LM's arrangement (Shoeybi et al., 2019) is worth memorising because everything since copies it. For the MLP block $Y = \mathrm{GeLU}(XA)B$: shard $A$ column-wise so each rank computes $\mathrm{GeLU}(XA_i)$ with no communication — GeLU is elementwise, so column sharding survives it — then shard $B$ row-wise so each rank produces a partial sum of $Y$, and finish with one all_reduce. Two matmuls, one collective. For attention, the head dimension shards naturally: each rank owns whole heads, computes them independently, and the output projection is row-parallel, giving one all_reduce per block. So a transformer layer costs 2 all-reduces forward and 2 backward, on activations of size $\text{batch} \times \text{seq} \times d_{\text{model}}$.

That activation traffic scales with batch and sequence length, unlike DDP's gradient traffic which scales with parameter count. It is the reason tensor parallel degree rarely exceeds the number of GPUs inside one node — typically 8.

Composing the three

Production 3-D parallelism arranges the world as $\text{world} = d \times t \times p$ (data, tensor, pipeline). The received ordering, from Narayanan et al.'s empirical study: tensor parallel within a node up to the NVLink domain; pipeline parallel across nodes, since point-to-point stage transfers are small and latency-tolerant; data parallel on the outside. FSDP can replace or compose with the data dimension through a multi-dimensional DeviceMesh. Sequence and context parallelism add a fourth axis for long contexts, sharding the sequence dimension of activations that tensor parallelism leaves replicated.

References

  • Huang et al. (2019), GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism — microbatching and the bubble formula.
  • Narayanan et al. (2019), PipeDream: Generalized Pipeline Parallelism for DNN Training — 1F1B and weight stashing.
  • Shoeybi et al. (2019), Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — the column/row transformer arrangement.
  • Narayanan et al. (2021), Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21 — interleaved schedules and the 3-D composition rules.
  • Qi et al. (2024), Zero Bubble Pipeline Parallelism — splitting the backward pass to remove the bubble.

What to learn next