Volume II · Splitting the Work
Cutting a Model Into Pieces
A model too large for one GPU's memory has to be split somehow, and there is more than one axis to split along. The one every student half-knows already is data parallelism — copy the whole model onto every GPU, show each copy different data, average the gradients. The one that actually cuts inside a layer is tensor parallelism, and the specific insight that makes it affordable — rather than a communication disaster — is the subject of this lecture.
The mental model
Splitting a matrix multiply across GPUs looks like it should need a synchronization after every single multiply. Megatron-LM's insight is to split the first and second matrices of a transformer block in matching, complementary ways, so the two multiplies compose without an intermediate sync — leaving exactly one all-reduce per block, not one per matrix.
The baseline: data parallelism
Data parallelism is the least novel and most load-bearing idea in distributed training: put an identical full copy of the model on every GPU, split a batch of training examples across those GPUs so each one sees a different slice of data, run the forward and backward pass independently on each copy, and then average the resulting gradients across all copies before applying an update — so every replica ends the step with the same weights again. Communication happens once per batch (an all-reduce of gradients), which is cheap relative to how much computation happens between those syncs. Its limitation is equally simple: every GPU needs to hold the entire model, so data parallelism alone cannot rescue a model that does not fit on one GPU in the first place. That is the problem tensor parallelism and, next lecture, pipeline parallelism exist to solve.
Megatron-LM: cutting inside the layer
Shoeybi et al.'s Megatron-LM (arXiv:1909.08053, 2019) takes the opposite cut: instead of replicating the whole model, split individual layers' weight matrices across GPUs, so each GPU holds only a slice of every layer. The paper's own numbers are worth stating plainly — implemented with a few inserted communication operations in native PyTorch, no custom compiler or new framework, it trained an 8.3-billion-parameter transformer across 512 GPUs, sustaining 15.1 PetaFLOPs at 76% scaling efficiency relative to a single-GPU baseline.
The mechanism is specific to how a transformer's MLP block is shaped, and it is
worth deriving rather than taking on faith. A transformer MLP block is two matrix
multiplies with a nonlinearity between them: Y = GeLU(X A), then
Z = Y B. Split matrix A by its columns across GPUs — each GPU
gets a vertical slice of A — and each GPU can independently compute its own slice of
Y with no communication at all, because column-splitting A means each GPU's slice of
Y depends only on the input X and its own slice of A. Now split the second matrix B
by its rows, matching the column-split of A: each GPU multiplies its slice
of Y by its corresponding row-slice of B, and the partial results across GPUs must
be summed to form the true Z — one all-reduce. Because A was split by columns and B
by rows in a matching way, the nonlinearity in between never needs to see a
complete, un-split Y — it can be applied to each GPU's partial slice independently.
The result: one all-reduce per MLP block, not one after the first multiply and
another after the second.
Why this stays inside one fast fabric
One all-reduce per transformer block is cheap only if the link doing that all-reduce is fast — and a transformer has many blocks per forward pass, so this communication happens repeatedly within a single training step. This is exactly the constraint Lecture 2's fabric hierarchy predicts: tensor parallelism is practical when its GPUs sit inside one NVLink-connected node, where an all-reduce is cheap relative to compute, and becomes impractical the moment it has to cross the much slower inter-node InfiniBand/Ethernet tier. That is why tensor parallelism is, in practice, kept inside a single node's worth of GPUs, while data parallelism's much rarer per-batch communication can tolerate crossing the slow inter-node fabric without much penalty.
Communication frequency is the real variable
Data parallelism and tensor parallelism are not "more" and "less" parallel in some abstract sense — they differ in how often they need to talk. Data parallelism syncs once per batch; tensor parallelism syncs many times per forward pass. Put the frequent communicator on the fast fabric, the infrequent one on whatever fabric is available.
Where this goes next
Narayanan et al.'s follow-up paper (arXiv:2104.04473, 2021) extends Megatron-LM with an interleaved pipeline schedule and formalizes combining tensor parallelism, pipeline parallelism (Lecture 4), and data parallelism into what the paper calls "3D parallelism" — the composition that production-scale training actually runs, not any one of the three dimensions alone.