Parallelism: data, tensor and pipeline (the ideas)

Step 1 Understand it

The problem it solves. Large models don't fit on one GPU, and even when they do, one GPU would take years to train them. Splitting the work is the only way to train at scale, and choosing how to split it — given the model size and how fast the GPUs can talk to each other — decides how fast and how cheaply a model trains.

Watch it happen

The idea

Data parallel: every GPU holds a full copy of the model, gets different data, and gradients are averaged. Tensor parallel: one layer's matrices are split across GPUs. Pipeline parallel: different layers live on different GPUs. FSDP and ZeRO split weights, gradients and optimiser states across GPUs to save memory.

How it works, step by step

1. Data parallel (DDP): every GPU holds the whole model and processes a different slice of the batch.
2. After backward, an all-reduce averages the gradients across GPUs, so every copy takes the same step and stays identical.
3. FSDP / ZeRO: like data parallel, but each GPU stores only a shard of the weights, gradients and optimiser states, gathering what it needs just in time. Much less memory per GPU.
4. Tensor parallel: split each large matrix across GPUs; each computes part of every layer, and they exchange results inside each layer. Needs very fast links, so it is used within one machine.
5. Pipeline parallel: put different layers on different GPUs and pass activations along, with micro-batches to keep stages busy. Idle time at the start and end is the "bubble".
6. Large runs combine them ("3-D parallelism"): tensor parallel within a machine, pipeline across machines, data parallel across groups.

Picture it

flowchart TB
  subgraph DP["Data parallel"]
    G1["GPU 1: whole model<br/>batch A"]
    G2["GPU 2: whole model<br/>batch B"]
  end
  subgraph TP["Tensor parallel"]
    T1["GPU 1: left half<br/>of every layer"]
    T2["GPU 2: right half<br/>of every layer"]
  end
  subgraph PP["Pipeline parallel"]
    P1["GPU 1: layers 1–12"] --> P2["GPU 2: layers 13–24"]
  end
  G1 <-->|"average gradients"| G2

Analogy

Cooking a huge dinner. Data parallel: several kitchens cook the full menu for different tables. Tensor parallel: several chefs cut one giant cake together. Pipeline: an assembly line, one station per course.

Where the analogy breaks

Separate kitchens don't need to talk. GPUs talk constantly, and the speed of that communication often decides which method wins.

Worked example

The model fits on one GPU and you want speed → data parallel. The model does not fit → FSDP, tensor or pipeline parallel.

Must-know points

• Model fits on one GPU → data parallel; doesn't fit → FSDP, tensor or pipeline parallel.
• Communication speed often decides what works.
• All-reduce keeps data-parallel copies in sync.
• The slowest GPU sets the pace for everyone.
• Pipeline bubbles shrink with more micro-batches.

Common mistakes

• Using tensor parallelism across slow network links.
• Forgetting that the effective batch size grows with the number of data-parallel GPUs, so the learning rate may need adjusting.
• Not checking that the multi-GPU run matches a single-GPU run on a small test.

Step 2 Go further (optional)

Optional. The lesson above is written to be enough on its own. If you want a second explanation, start with the one marked Worth it.

• Read

  The Ultra-Scale Playbook

  Hugging Face · interactive diagrams

  Worth it: A long but beautifully illustrated guide to every kind of parallelism, with interactive memory and speed calculators, from people who train large models.

• Read

  Getting Started with Distributed Data Parallel

  PyTorch tutorials

• Paper

  Megatron-LM (2019)

  Shoeybi et al. · arXiv · optional

Step 3 Practise

Try each exercise yourself first: on paper, out loud or in code. Then open the answer and compare it with yours.

1. ExplainSay it in plain words, with no notes.

   Explain data, tensor and pipeline parallelism with the dinner-kitchen picture.

   Show answer

   Cooking a huge dinner for a wedding.

   Data parallel: several kitchens each cook the full menu, each for different tables, and at the end of every course the head chefs compare notes so every kitchen improves its recipe in the same way. It needs each kitchen to be big enough for the whole menu.

   Tensor parallel: one cake is too big for any single chef, so several chefs work on the same cake at the same time, each on their own slice, constantly passing tools and checking with each other. They must stand close together (fast links in one machine).

   Pipeline parallel: an assembly line, with one station per course. Starters go to station 1 and then move on, while station 1 starts the next table's starters, so all stations stay busy, except at the very start and end of the service.

2. DeriveWrite the maths, or the steps, from memory.

   Draw data parallelism for 8 GPUs and a model that fits on one. Then draw tensor and pipeline parallelism for a model that needs 4 GPUs.

   Show answer

   Data parallelism, 8 GPUs, a model that fits on one: every GPU holds a full copy of the model. Each takes a different slice of the batch (8 slices), runs forward and backward on its own, and then all 8 average their gradients with an all-reduce. Every copy applies the same averaged update, so they stay identical. It scales the batch, not the model.

   Tensor parallelism, a model that needs 4 GPUs: each big weight matrix is split across the 4 GPUs, for example by columns. For one layer, every GPU multiplies the input by its quarter of the matrix, then they exchange results (an all-reduce or all-gather) before the next layer. Very chatty, so it is used inside one machine with fast links.

   Pipeline parallelism, 4 GPUs: the layers are split into 4 consecutive stages, layers 1 to 6 on GPU 0, 7 to 12 on GPU 1, and so on. Activations flow from stage to stage. To keep every GPU busy, the batch is cut into micro-batches that follow each other through the pipeline; the idle time at the start and end is the "bubble".

   Large models combine them: tensor parallel within a machine, pipeline across machines, data parallel across groups, often with the optimiser state sharded too (FSDP or ZeRO).

   textCopy

   DATA PARALLEL (model fits on 1 GPU)        TENSOR PARALLEL (1 layer split 4 ways)
   GPU0: full model, batch slice 1 ┐          input ─┬▶ GPU0: W[:, 0:¼]  ┐
   GPU1: full model, batch slice 2 ├ all-reduce      ├▶ GPU1: W[:, ¼:½]  ├ gather ▶ next layer
   ...                             │ gradients       ├▶ GPU2: W[:, ½:¾]  │
   GPU7: full model, batch slice 8 ┘                 └▶ GPU3: W[:, ¾:1]  ┘
   
   PIPELINE PARALLEL (layers split into 4 stages, micro-batches m1, m2, m3 …)
   GPU0 [layers 1–6]   m1 m2 m3 m4
   GPU1 [layers 7–12]     m1 m2 m3 m4
   GPU2 [layers 13–18]       m1 m2 m3 m4
   GPU3 [layers 19–24]          m1 m2 m3 m4     (the empty corners are the "bubble")

3. BuildCode it from scratch in NumPy, with no ML library.

   Run data-parallel training with PyTorch DDP on 2 GPUs, or 2 CPU processes. Check the gradients match single-device training.

   Show answer

   DistributedDataParallel with 2 CPU processes and the gloo backend, so no GPU is needed. Each process builds the same model (same seed), takes half of the batch, and after backward DDP averages the gradients across the processes automatically. Then compare with one process computing the gradient of the whole batch: they match to better than 1e-8. That equality is the guarantee data parallelism gives you: N workers on N slices compute exactly the update one worker would compute on the full batch.

   On real GPUs, change the backend to "nccl", put each process on its own device, and launch with torchrun --nproc_per_node=2 train.py.

   pythonCopy

   import os
   import torch
   import torch.distributed as dist
   import torch.multiprocessing as mp
   from torch.nn.parallel import DistributedDataParallel as DDP
   
   def make_model():
       torch.manual_seed(0)
       return torch.nn.Sequential(torch.nn.Linear(10, 32), torch.nn.ReLU(), torch.nn.Linear(32, 1))
   
   def worker(rank, world, X, Y, queue):
       os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29511")
       dist.init_process_group("gloo", rank=rank, world_size=world)
       model = DDP(make_model())
       xs, ys = X.chunk(world)[rank], Y.chunk(world)[rank]            # this worker's slice
       torch.nn.functional.mse_loss(model(xs), ys).backward()        # DDP averages gradients here
       if rank == 0:
           queue.put(model.module[0].weight.grad.clone())
       dist.destroy_process_group()
   
   if __name__ == "__main__":
       torch.manual_seed(1)
       X, Y = torch.randn(64, 10), torch.randn(64, 1)
       queue = mp.get_context("spawn").SimpleQueue()
       mp.spawn(worker, args=(2, X, Y, queue), nprocs=2)
       ddp_grad = queue.get()
       single = make_model()
       torch.nn.functional.mse_loss(single(X), Y).backward()
       print("max difference from single-device training:", (ddp_grad - single[0].weight.grad).abs().max().item())

4. BreakShow when it fails, and why.

   Make one process slower on purpose. What happens to the speed of the whole group, and why?

   Show answer

   Make one process sleep for a while in every step. Now every step of the whole group takes as long as the slowest process: in the run below, both workers' step time jumps to roughly the sleeper's, even though the other worker has nothing extra to do.

   Why: the gradient all-reduce is a synchronisation point. No worker can finish the step until every worker has contributed its gradients, so the fast one waits at the all-reduce, idle. With 1,000 GPUs, one slow GPU, a bad network link, or one worker that got an unusually long batch slows all 1,000 down.

   That is why large training runs watch for stragglers, keep batches the same size on every worker, overlap communication with computation, and replace failing machines quickly.

   pythonCopy

   import os
   import time
   import torch
   import torch.distributed as dist
   import torch.multiprocessing as mp
   from torch.nn.parallel import DistributedDataParallel as DDP
   
   def worker(rank, world, slow_rank, queue):
       os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29512")
       dist.init_process_group("gloo", rank=rank, world_size=world)
       torch.manual_seed(0)
       model = DDP(torch.nn.Linear(100, 100))
       x = torch.randn(32, 100)
       times = []
       for _ in range(5):
           start = time.perf_counter()
           if rank == slow_rank:
               time.sleep(0.2)                                    # this worker is slow
           model(x).sum().backward()                              # all-reduce waits for everyone
           times.append(time.perf_counter() - start)
       queue.put((rank, slow_rank, sum(times) / len(times)))
       dist.destroy_process_group()
   
   if __name__ == "__main__":
       ctx = mp.get_context("spawn")
       for slow in [None, 1]:
           queue = ctx.SimpleQueue()
           mp.spawn(worker, args=(2, slow, queue), nprocs=2)
           results = sorted(queue.get() for _ in range(2))
           print(f"slow worker = {slow}: " + ", ".join(f"worker {r} step {t * 1000:.0f} ms" for r, _, t in results))

Step 4 Prove it

Check you understood

Tick each check when you can do it without notes