Chapter 12 of 16
Parallelism II: model
Tensor, pipeline, sequence and expert parallelism; the pipeline bubble.
Slides for this chapter: Transformer Training 11, Model Parallelism
Data parallelism gives every device the whole model. When the model is too big for one device, it has to be cut up: across devices layer by layer (a pipeline), inside each layer (tensor parallelism), along the sequence, or among experts. Below, a batch flows through a four-device pipeline in micro-batches, every pass simulated as an event; the empty space is the bubble, time the devices spend waiting.
Loading the animation…
Concept
Pipeline parallelism. Give each of devices a contiguous block of layers, chapter 4's stack cut into stages: device 1 holds the first quarter of the layers, device 4 the last. A sequence's activations flow forward from device to device and its gradients flow back. One batch at a time would leave three devices idle at any moment, so the batch is cut into micro-batches that follow each other through the stages. GPipe (Huang et al. 2019) runs every micro-batch forward, then every one backward; at the start and end of each step the pipeline fills and drains, and those idle slots are the bubble, a fraction of the useful work: 37.5% for 4 devices and 8 micro-batches in the simulation (the formula gives the same, to the last bit; results.md §12). Measured against the whole step instead (the useful work plus the bubble, slots per device), the same idle time is , here: the form Serving, chapter 9 and Trade-offs, chapter 8 use. Both describe the same schedule. 1F1B (one forward, one backward) interleaves each device's forwards and backwards once the pipeline is full: the bubble is the same, but a device holds at most micro-batches' activations instead of all (Narayanan et al. 2021, §2.2). Megatron-LM's interleaved schedule gives each device smaller chunks of layers spread round the pipeline (device 1 holds layers 1–2 and 9–10, say), cutting the bubble to : 18.8% with , 9.4% with , at the price of times as many sends between devices. Pick a schedule above and watch where the idle time goes.
Tensor parallelism. Cut each layer's matrices instead (Shoeybi et al. 2019, Megatron-LM). The feed-forward network's first matrix (Decoder, chapter 4) is split by columns, so each of devices computes its own slice of the hidden layer, GELU included, with no communication; the second matrix is split by rows, so each device produces a partial sum of the output, and one all-reduce adds them. Attention splits the same way, a group of heads per device (Decoder, chapter 3: the heads are already independent). That is two all-reduces of the activations per layer in the forward pass and two in the backward pass, every layer, every micro-batch: far more communication than data parallelism's once-per-step gradient sync, which is why tensor parallelism stays inside one server, on NVLink. In the systems model, Megatron's 18.4-billion-parameter run (, no pipeline) spends 1.67 s of its 10.61 s step in those all-reduces.
Sequence and context parallelism. Tensor parallelism shares the matrix products but not LayerNorm or dropout, whose activations each device would still hold in full. Sequence parallelism (Korthikanti et al. 2022) splits those along the sequence, turning each all-reduce into a reduce-scatter and an all-gather (the same bytes) and cutting the activations per device by . For very long sequences, context parallelism splits the sequence itself across devices for attention too: each device keeps its own queries and gets everyone's keys and values, passed round a ring while computing (Liu et al. 2023, ring attention) or all-gathered first, as Llama 3 does for its 128K-token stage (Grattafiori et al. 2024, §3.3.2).
Expert parallelism. A mixture-of-experts model replaces each feed-forward network with many experts, of which each token uses a few: one or two in the early designs (Lepikhin et al. 2020, GShard; Fedus et al. 2021, Switch Transformers), typically 2 to 8 in current models, often with a shared expert every token uses (DeepSeek-V3 routes each token to 8 of 256 experts plus 1 shared; Architectures, chapter 4). The experts live on different devices; every layer sends each token to its expert's device and back with two all-to-all exchanges, where every device sends a different piece to every other.
3-D and 4-D layouts. Real runs combine them: tensor parallelism inside a server, pipeline stages across servers, data parallelism (often sharded, ZeRO or FSDP, chapter 11) outermost. Llama 3 405B trained with four dimensions, [tensor, context, pipeline, data], on up to 16,384 GPUs; its first stage used 8-way tensor, 16-stage pipeline and 64-way data parallelism (Table 4). Choosing a layout is a trade between the memory each device must hold and the time lost to communication and bubbles. The explorer below runs this site's systems model on three published model shapes: try fitting GPT-3 into 80 GB, then see what the bubble costs with few micro-batches.
Loading the animation…
How good is the systems model? It is arithmetic with one fitted constant: the fraction of peak the matrix products reach, 0.473, fitted on Megatron-LM's smallest run (1.7B parameters on 32 A100s, no tensor or pipeline parallelism). Every other row of the chart is a prediction for a published run: nine more Megatron-LM runs up to a trillion parameters on 3,072 A100s (Narayanan et al. 2021, Table 1) and the three stages of Llama 3 405B on H100s (Grattafiori et al. 2024, Table 4), with only the shapes, layouts, batch sizes and GPU specifications as inputs. **Its stated band: every prediction within 18.6% of the published throughput** (Megatron rows within 18.6%, Llama 3 rows within 10.3%; mean error 10.8% ). Every prediction is low: the model overlaps no communication with compute, and one cannot follow matrix products getting more efficient as models widen (Megatron's measured throughput per GPU rises from 137 to 163 TFLOP/s across its table; the model's stays near 130). Where the papers are silent the model assumes one sequence per micro-batch and the most interleaving the layer count allows; results.md §11a lists every assumption.
The hardware underneath. Racks, chapter 9: training clusters places a tensor × pipeline × data parallel job on NVL72 racks and on 8-GPU servers, group by group, and counts the pipeline hops that cross the network; chapter 4 is the NVLink domain that tensor parallelism stays inside (the Racks site, in the Hardware track).
Maths
Tensor parallelism's communication, per layer per micro-batch of sequences of length and width over devices: 4 all-reduces of bytes (6 with full recomputation), each costing
Pipeline, stages, micro-batches, chunks per device, a micro-batch's forward and backward on one stage taking :
Activations on the first stage (Korthikanti et al., Table 2 and §4.2.3), bytes per layer with tensor and sequence parallelism, times layers' worth (the first stage holds micro-batches of its layers) and when interleaved:
Step time in the systems model: , with the compute of one micro-batch on one stage at of peak, its tensor-parallel all-reduces, its context-parallel all-gathers, and nothing overlapped.
Code
The 1F1B order of one device's work, from the pipeline simulation in
src/lib/training/systems.ts (matched to scripts/training/systems.py):
if (kind === "1f1b") {
const warm = Math.min(p - s - 1, m);
for (let i = 0; i < warm; i++) out.push(["F", i, 0]);
let f = warm;
let b = 0;
while (f < m) {
out.push(["F", f++, 0]);
out.push(["B", b++, 0]);
}
for (let i = b; i < m; i++) out.push(["B", i, 0]);
return out;
}
Comments
Be the first to leave a comment on this section.
Sign in (top-right) to leave a comment.