LLMUnit 5 · Generative models & scaling · Day 24

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Cut a model too big for one chip into stages on several chips, and split each batch into micro-batches so the stages work at the same time instead of waiting for each other.

AuthorsYanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, Zhifeng Chen (Google)
Year / venue2019 · NeurIPS (arXiv Nov 2018)
Core ideaMicro-batch pipeline parallelism with synchronous updates
Study time2–3 h

Original: arXiv:1811.06965 (later versions are titled "GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism") · Code, notebook and exercises: 30u30 Day 24

Start here (no maths needed)

When a neural network is too big for one computer chip, you can split it across several chips and keep them all busy, like an assembly line.

Think of it like this

Think of a sandwich assembly line with four people: one adds bread, the next adds cheese, the next adds salad, the last wraps it. If you make one giant order all at once, three people stand around waiting while one works. But if you split the order into many small batches, everyone is busy at the same time. There's only a little waiting at the very start and end. GPipe does this with computer chips.

How it works

  1. Cut the network into stages, one stage per chip.
  2. Split each big batch of training examples into small batches, called micro-batches.
  3. Send the micro-batches through the chips one after another, so all chips work at once.
  4. Each chip remembers only a little and redoes some work later. This saves memory.
  5. Add up what was learned from all the micro-batches and update the network once. The result is the same as using one giant chip.

Why robots care

The big AI models that help robots understand instructions and pictures are far too large for one chip. Splitting them across many chips like this is part of how they get trained.

Try it

Words to know

Chip (accelerator)
A piece of computer hardware that runs the network fast.
Stage
One part of the network, running on one chip.
Micro-batch
A small group of training examples sent through the line together.
Bubble
Time when a chip is idle, waiting for work.
Pipeline
Work flowing through stages one after another, like an assembly line.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Summary

TL;DR

If a network is a sequence of layers, you can put the first few layers on chip 1, the next few on chip 2, and so on. Done naively, only one chip works at a time, because each one waits for the previous chip's output. GPipe splits each mini-batch into M micro-batches and streams them through the K stages like an assembly line. Gradients from all micro-batches are added up and applied once, so the update is exactly the same as training on the full batch on one huge device. Add re-materialisation (recompute activations in the backward pass instead of storing them) and memory use drops too. With this, Google trained a 557M-parameter image model to 84.4% ImageNet top-1, and a 6B-parameter, 128-layer translation Transformer covering 103 languages.

2. The core idea

Write the network as L layers and group consecutive layers into K cells, one per accelerator. Cell k computes F_k forward and B_k backward. Split a mini-batch of size N into M micro-batches:

forward:  micro-batch m runs F_1 → F_2 → … → F_K
backward: micro-batch m runs B_K → … → B_1
update:   g = Σ_{m=1..M} g_m   applied once, after all M are done

bubble fraction ≈ (K − 1) / (M + K − 1)
peak activation memory ≈ O(N + (L/K) · (N/M))   with re-materialisation
                 vs      O(N · L)               without

Interactive: the pipeline schedule

Rows are accelerators, columns are time steps. Forward cells are solid, backward cells are lighter, and blank cells are the bubble. Raise M and watch the idle fraction shrink.

Idle fraction (counted)
–
(K−1)/(M+K−1)
–
M / K
–

3. Deep dive

Three ways to split training

Large modern training runs combine all three, plus expert parallelism for mixture-of-experts models. Module 05 shows routing among hundreds of experts spread across GPUs, the newest axis on this list.

Why synchronous matters

PipeDream, an earlier pipeline system, overlapped forward and backward passes of different batches. That kept chips busier but meant gradients were computed against slightly stale weights, and it had to store several weight versions. GPipe accepts some bubble in exchange for exact, synchronous updates. Training gives the same result whether you use 2 or 16 stages, which makes debugging and reproducing far easier.

What the numbers show

On TPUs with M = 32, going from 2 to 8 stages raised Transformer throughput from 1.8 to 6.3 (normalised), about 3.5× for 4× the chips. AmoebaNet scaled less well because its layers have uneven cost, so the stages couldn't be balanced. With M = 1 there's no pipelining at all, and adding chips barely helped. On GPUs connected only through PCIe, 2 → 8 stages still gave 2.7× (AmoebaNet) and 3.3× (Transformer). The translation experiments also found that, at equal size, a deeper model beat a wider one by large margins on low-resource languages. That's a result GPipe made affordable to test, and it feeds into the scaling story of Day 25.

Gotchas

4. Why it matters for robotics

5. How to study it

  1. Read sections 2 and 3 of the paper (about 1 h). Redraw Figure 2(c) by hand for K = 3, M = 4 and count the idle cells.
  2. In the 30u30 Day 24 folder, read implementation.py: find the micro-batch split, the schedule loop, gradient accumulation and checkpointing.
  3. Run train_minimal.py with different --partitions and --micro-batches values and compare against the bubble formula.
  4. Do the exercises: micro-batch splitting, the pipeline forward pass, gradient accumulation, re-materialisation, then the full wrapper.
  5. Stretch: check that the accumulated gradient over M micro-batches equals the full-batch gradient, to numerical precision, for a model without batch norm.

6. Self-check

What is the pipeline bubble and how do you shrink it?

Idle time while the pipeline fills and drains, about (K−1)/(M+K−1) of the total. More micro-batches per stage shrinks it. GPipe found M ≥ 4K enough to make it negligible.

Why are GPipe's gradients identical to single-device training?

All micro-batches use the same weights, and their gradients are summed and applied once per mini-batch. Nothing is stale.

What does re-materialisation trade?

Compute for memory. Each stage stores only boundary activations and recomputes the rest during the backward pass, costing roughly one extra forward pass.

Why did the Transformer scale more linearly than AmoebaNet?

Transformer layers all cost about the same, so stages balance well. AmoebaNet's layers vary, so one stage is slower and holds the others back.

When is pipeline parallelism not enough?

When a single layer doesn't fit on one device, or when there are too few micro-batches to hide the bubble.

7. Further reading