When one device cannot hold the model, cut it up: layer by layer (a pipeline, with its bubble), inside each layer (tensor parallelism), along the sequence, or among experts. A live pipeline simulation, 3-D and 4-D layouts, and a systems model checked against published runs.
This deck follows chapter 12, parallelism II: model, with a pipeline animation and a layout explorer. Numbers: results.md §11a and §12 (the site’s training-systems model and its discrete-event pipeline simulation).
Data parallelism gives every device the whole model. When the model is too big for one device (deck 09), it has to be cut up:
Each device holds a contiguous block of layers; activations flow forward from device to device and gradients flow back.
Each layer’s matrices are split across devices, which combine partial results with collectives.
Activations, and for long contexts attention itself, are split by position.
A mixture-of-experts model puts different experts on different devices.
Give each of \(p\) devices a quarter (say) of the layers. One batch at a time would leave all but one device idle, so the batch is cut into \(m\) micro-batches that follow each other through the stages. At the start and end of each step the pipeline fills and drains; those idle slots are the bubble.
$$T = \Big(m + \tfrac{p-1}{v}\Big)(t_f + t_b), \qquad \text{bubble} = \frac{p-1}{v\,m}$$Every micro-batch forward, then every one backward (Huang et al. 2019). Each device must hold all \(m\) micro-batches’ activations.
One forward, one backward: once the pipeline is full, each device alternates. The same bubble, but a device holds at most \(p\) micro-batches’ activations (Narayanan et al. 2021, §2.2).
A discrete-event simulation of every forward (1 time unit) and backward (2 units) on every device, each waiting for its inputs. Choose the schedule and sizes; the idle time is the bubble.
Megatron-LM’s interleaved schedule gives each device \(v\) smaller chunks of layers spread round the pipeline (device 1 holds layers 1–2 and 9–10, say), cutting the bubble to \((p-1)/(vm)\), at the price of \(v\) times as many sends between devices. The site’s simulation against the closed forms (results.md §12):
| Schedule | Devices p | Micro-batches m | Chunks v | Simulated bubble | Closed form |
|---|---|---|---|---|---|
| gpipe | 4 | 8 | 1 | 0.3750 | 0.3750 |
| 1f1b | 4 | 8 | 1 | 0.3750 | 0.3750 |
| interleaved | 4 | 8 | 2 | 0.1875 | 0.1875 |
| interleaved | 4 | 8 | 4 | 0.0938 | 0.0938 |
| 1f1b | 8 | 32 | 1 | 0.2188 | 0.2188 |
| interleaved | 8 | 32 | 4 | 0.0547 | 0.0547 |
| gpipe | 16 | 32 | 1 | 0.4688 | 0.4688 |
| 1f1b | 16 | 64 | 1 | 0.2344 | 0.2344 |
Largest difference between simulation and formula: 0.0.
Cut each layer’s matrices instead (Shoeybi et al. 2019, Megatron-LM):
The FFN’s first matrix (Decoder, chapter 4) split by columns lets each of \(t\) devices compute its own slice of the hidden layer, GELU included, with no communication; the second, split by rows, gives each a partial sum of the output, and one all-reduce adds them. Attention splits the same way, a group of heads per device (the heads are already independent, Decoder, chapter 3).
Two all-reduces of the activations per layer in the forward pass and two in the backward, every layer, every micro-batch: far more traffic than data parallelism’s once-per-step gradient sync, so tensor parallelism lives on NVLink. In the systems model, Megatron’s 18.4-billion-parameter run (\(t = 8\), no pipeline) spends 1.67 s of its 10.61 s step in those all-reduces.
Tensor parallelism shares the matrix products but not LayerNorm or dropout, whose activations each device would still hold in full. Sequence parallelism splits those along the sequence, turning each all-reduce into a reduce-scatter and an all-gather (the same bytes) and cutting activations per device by \(t\) (Korthikanti et al. 2022).
For very long sequences, split the sequence itself 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 (§3.3.2).
Mixture-of-experts layers (GShard; Switch Transformers) send each token to its expert’s device and back with two all-to-all exchanges per layer, where every device sends a different piece to every other.
Real runs combine them: tensor parallelism inside a server, pipeline stages across servers, data parallelism (often sharded, ZeRO or FSDP) 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 (Grattafiori et al. 2024, Table 4).
8 × 16 × 64 = 8,192 GPUs. Choosing a layout trades the memory each device must hold against the time lost to communication and bubbles; the site’s layout explorer runs its systems model on published shapes.
Arithmetic with one fitted constant: the fraction of peak the matrix products reach, η = 0.473, fitted on Megatron-LM’s smallest run. Every other row is a prediction for a published run, with only shapes, layouts, batch sizes and GPU specifications as inputs (Narayanan et al. 2021, Table 1; Llama 3, Table 4; results.md §11a):
Stated band: every prediction within 18.6% of the published throughput (Megatron 18.6%, Llama 3 10.3%; mean 10.8%). Every prediction is low: the model overlaps nothing, 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). Assumptions where the papers are silent: one sequence per micro-batch, the most interleaving the layer count allows, BF16 on the wire.
Animated: chapter 12, parallelism II: model · hardware: NVIDIA GPU Architectures, Google TPUs · next: deck 12, training over slow links