Splitting the Model

It is late 2018. Transformer-based language models are about to cross a threshold where the model itself — not the training set — no longer fits in one accelerator’s memory, even at a batch size of one. Data parallelism, the subject of the previous chapter, has no answer to that problem: it replicates the whole model on every device. The only way forward is to cut the model itself into pieces and hand the pieces to different chips.

Where We Were

Chapter 8 covers how ring all-reduce and the linear scaling rule let a team throw more GPUs at a training run by splitting the batch across replicas, each holding a full copy of the model. That trick has a hard ceiling: every replica still needs enough device memory to hold its own parameters, gradients, and optimizer state, plus the activations for its slice of the batch. Chapter 4 covers why that ceiling is a real, physical one — accelerator memory grows far slower than the compute attached to it, and by the late 2010s a single high-end GPU carried on the order of tens of gigabytes.

The field had already brushed up against this limit once before. AlexNet, split across two GTX 580 cards in 2012 because 60 million parameters and their activations did not fit in 3 GB of memory on one card, is covered in Chapter 2. That split was hand-coded for one network on one pair of cards, with cross-GPU communication restricted to a few specific layers to keep overhead down. It solved the immediate problem but generalized to nothing: every new model needing a manual split needed its own bespoke wiring, decided by hand for that specific architecture and that specific pair of cards.

By 2018, two things had changed. First, the Transformer — covered in language models, Chapter 5 — had become the default architecture for sequence modeling, and it scales by stacking more identical layers and widening the same few matrix multiplications — a much more regular structure to split than AlexNet’s mix of convolutions and fully connected layers. Second, model sizes were about to grow past what any single accelerator could hold, first in machine translation, then in language modeling: BERT, at 340 million parameters, shipped in October 2018 (language models, Chapter 6), and GPT-2 followed in February 2019 at 1.5 billion (language models, Chapter 7). The field needed a general, reusable way to split a model across devices — not a one-off hack per architecture.

Two different places to make the cut suggested themselves, because a neural network is, structurally, both a sequence of layers and, within each layer, a set of large matrix multiplications. Cutting between layers means each device holds a handful of complete layers and passes activations to the next device in sequence — conceptually close to what a human reading a diagram of a deep network would draw first, since the layers are already drawn as discrete boxes in a line. Cutting within a layer means every device holds a slice of every layer, and the devices jointly compute each layer’s output before moving on together — a less obvious cut, because it requires reasoning about the linear-algebra structure of a specific operation rather than just the network’s block diagram. Both cuts turned out to matter, and for different reasons: one keeps communication rare but forces devices to sit idle waiting for work to arrive, and the other keeps devices busy in lockstep but requires communication after nearly every operation.

The Breakthrough

Two independent answers arrived within eleven months of each other, and they split the model along two different axes.

Huang et al., at Google, published GPipe in November 2018 (Huang et al., 2018). GPipe partitions a network’s layers into $K$ consecutive stages and places each stage on its own accelerator — the model is cut between layers. The obvious way to run this, one microbatch at a time through the full stage sequence, leaves every device but one idle at any given moment: while stage 1 processes microbatch 1, stages 2 through $K$ have nothing to do yet. GPipe’s fix is to split each minibatch into $M$ smaller microbatches and pipeline them through the stages, so that stage 1 starts on microbatch 2 as soon as it hands microbatch 1 to stage 2, and so on — the same fill-and-drain idea as an assembly line. GPipe also recomputes activations during the backward pass rather than storing all of them, an application of the gradient checkpointing technique covered in Chapter 10, which lets each stage hold far less in memory at the cost of some repeated forward compute. GPipe reported training a 557-million-parameter AmoebaNet-B image classifier split across 4 accelerators to 84.4% top-1 accuracy on ImageNet, and a 6-billion-parameter, 128-layer multilingual Transformer across up to 16 accelerators for translation across over 100 languages, outperforming every bilingual baseline.

Shoeybi et al., at NVIDIA, published Megatron-LM in September 2019 (Shoeybi et al., 2019), cutting the model along the other axis: within a layer rather than between layers. Rather than sending a whole matrix multiplication to one device, Megatron-LM splits the weight matrices of each transformer block’s MLP and self-attention sublayers across devices, so that every device computes a slice of the same layer’s arithmetic in parallel and the slices are stitched back together with a small number of synchronization points — detailed under The Mathematics below. The same column-and-row split applies to multi-head self-attention almost for free, because the attention heads are already independent of each other: each device is simply assigned a subset of the heads, computes attention for those heads on its local slice of the query, key, and value projections, and the results are concatenated and reduced the same way as the MLP block. Because this intra-layer split needs communication after nearly every matrix multiplication, it depends on very fast device-to-device links — NVLink inside a server, not the slower InfiniBand fabric between servers — and is normally kept within a single multi-GPU node rather than spread across a network. Megatron-LM trained transformer models up to 8.3 billion parameters across 512 GPUs, sustaining 15.1 aggregate petaFLOP/s at 76% scaling efficiency relative to a single-GPU baseline that itself sustained 39 teraFLOP/s (30% of that GPU’s peak), and reported a WikiText-103 perplexity of 10.8 against a prior state of the art of 15.8, a LAMBADA accuracy of 66.5% against 63.2%, and a RACE accuracy of 90.9% against 89.4%.

The two techniques are complementary rather than competing: pipeline parallelism scales well across the slower links between servers because it only communicates activations at stage boundaries, while tensor parallelism scales well across the fast links inside a server because it communicates far more often but in smaller amounts. Smith et al. combined all three splitting strategies — tensor, pipeline, and data — into what they called 3D parallelism to train Megatron-Turing NLG, a 530-billion-parameter model, in January 2022 (Smith et al., 2022): 8-way tensor parallelism inside each server, 35-way pipeline parallelism across servers, and data parallelism layered on top using the ZeRO optimizer (covered in Chapter 10), spread across 560 DGX A100 servers — 4,480 NVIDIA A100 GPUs in total, sustaining 126 teraFLOP/s per GPU and a peak cluster throughput of roughly 1.4 exaFLOP/s at 16-bit precision.

The Mathematics

GPipe’s central cost is idle time. With $K$ pipeline stages and a minibatch split into $M$ microbatches, every stage must wait for the pipeline to fill before it has continuous work and wait again for it to drain at the end, giving a bubble — idle time with no useful computation — of:

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

Megatron-LM’s central move is splitting a single matrix multiplication so that most of it needs no communication at all. Inside a transformer MLP block, the first weight matrix is split by columns and the second by rows across $P$ devices, and each device computes:

$$Z = \sum_{i=1}^{P} \mathrm{GeLU}(X A_i)\, B_i$$

Symbol	Meaning
$X$	input activations, replicated on all $P$ devices
$A_i$	$i$-th column-slice of the first weight matrix, local to device $i$
$B_i$	$i$-th row-slice of the second weight matrix, local to device $i$
$P$	tensor-parallel degree, number of devices splitting this layer
$Z$	this layer’s output activations, after one all-reduce

3D parallelism is bookkeeping on top of both mechanisms: the total number of accelerators used is the product of the three parallel degrees chosen along each axis.

$$N_{\mathrm{GPU}} = d_{\mathrm{TP}} \times d_{\mathrm{PP}} \times d_{\mathrm{DP}}$$

Symbol	Meaning
$d_{\mathrm{TP}}$	tensor-parallel degree, split within a layer
$d_{\mathrm{PP}}$	pipeline-parallel degree, split across layers
$d_{\mathrm{DP}}$	data-parallel degree, split across the batch, Chapter 8
$N_{\mathrm{GPU}}$	total accelerators required by the run

What It Enabled

	Data parallelism only (Chapter 8)	+ Pipeline and tensor parallelism (this chapter)
What limits maximum model size	one device’s memory, no matter how many devices are added	the sum of memory across every device on every axis
Representative model	any model that fits on one accelerator	8.3B-parameter Megatron-LM model (2019); 530B-parameter Megatron-Turing NLG (2022)
Representative cluster	tens of GPUs, replicated	512 GPUs (Megatron-LM); 4,480 A100 GPUs (Megatron-Turing NLG)
Reported efficiency	near-linear with ring all-reduce, but model size unchanged	76% scaling efficiency at model-parallel size 8 (Megatron-LM); 126 teraFLOP/s per GPU, ~1.4 exaFLOP/s peak cluster throughput (Megatron-Turing NLG)

Parallelism axis	What it splits	Communication pattern	Owning chapter
Data parallelism	the batch, across full model replicas	one all-reduce of gradients per step	Chapter 8
Pipeline parallelism	the model’s layers, across devices	activations forward and gradients backward, only at stage boundaries	this chapter, GPipe
Tensor parallelism	individual matrix multiplications, within a layer	roughly two all-reduces per transformer layer	this chapter, Megatron-LM
ZeRO / FSDP sharding	optimizer state, gradients, and parameters, across replicas	gather and scatter of shards, layered on data parallelism	Chapter 10

The two techniques also carry different failure modes, which is why real training runs pick both rather than either. Pipeline parallelism’s bubble, from the equation above, means a poorly chosen microbatch count wastes accelerator-hours doing nothing; tensor parallelism’s frequent all-reduces mean a poorly chosen placement — splitting a layer across devices that only talk to each other over a slow network link — turns every forward pass into a communication-bound crawl. Megatron-Turing NLG’s 8-way tensor split inside a node and 35-way pipeline split across nodes is a direct answer to that constraint: put the chatty axis where the wire is fast, and the quiet axis where it is slow.

Splitting the model along two axes, and later stacking a third on top, turned “the model does not fit on one GPU” from a dead end into a cluster-configuration decision. Once tensor and pipeline parallelism existed as reusable libraries rather than per-model hacks, the practical bottleneck moved from “can we fit it” to “can we operate it” — keeping a training run with thousands of GPUs healthy for weeks, the subject of Chapter 11. It also opened a fourth axis that this chapter does not cover: instead of splitting a dense layer’s arithmetic across devices, a model can activate only a subset of a much larger set of per-device expert layers, an idea covered in Chapter 13. And the scaling behavior this infrastructure made possible — bigger models, trained on more GPUs, reliably improving on benchmarks — is what let researchers test scaling laws empirically rather than in theory, covered in language models, Chapter 10.

Key Papers

Paper	Authors	Year	Link	Why it matters
GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism	Huang et al.	2018	arXiv:1811.06965	First general, reusable library for splitting a network’s layers across accelerators and pipelining microbatches to keep them busy
Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism	Shoeybi et al.	2019	arXiv:1909.08053	Introduced intra-layer tensor parallelism for transformers, splitting individual matrix multiplications across GPUs with minimal communication
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B	Smith et al.	2022	arXiv:2201.11990	Combined tensor, pipeline, and ZeRO-based data parallelism into 3D parallelism to train a 530-billion-parameter model across 4,480 GPUs

PM Lens