Megatron-LM, annotated

2.2 Transformer language models and multi-head attention · original

Everyday picture

The paper works with both families of transformer: GPT-2, a decoder that writes left to right, and BERT, an encoder that reads in both directions. Each layer is the same two steps, self-attention (tokens exchange information) then a two-layer feed-forward network (each token alone), each wrapped in a residual connection. Both use GELU, and both normalize the input to each step rather than its output as the original transformer did.

Tiny example

GELU (the paper writes GeLU) is a smooth ReLU: x times the chance that a standard bell-curve variable falls below x. GELU(1) = 1 × 0.841 = 0.841; GELU(−1) = −1 × 0.159 = −0.159; GELU(2) = 1.954. Big positive inputs pass almost unchanged, big negative ones go to almost zero.

In Python:

import math
def gelu(x):
    # x times Φ(x), the bell curve's cumulative probability
    return x * 0.5 * (1 + math.erf(x / math.sqrt(2)))
round(gelu(1), 3), round(gelu(-1), 3), round(gelu(2), 3)  # → (0.841, -0.159, 1.954)

Why it matters

That GELU is not linear, GELU(a + b) ≠ GELU(a) + GELU(b), is the fact §3's whole design turns on. The transformer lesson builds the layer (TransformerBlock, with gelu).

2.3 Data and model parallelism in deep learning · original

Everyday picture

Data parallelism gives every GPU the whole model and a slice of the batch; growing the batch with the number of GPUs is weak scaling. It has a hard limit: the model must fit on one GPU. Model parallelism removes the limit, in one of two ways: pipeline parallelism (whole layers per GPU, fed by micro-batches, with idle bubbles) or distributed tensor computation (each operation split across GPUs), which is this paper's choice.

Tiny example

A 1.2-billion-parameter model trained with Adam in mixed precision needs about 16 bytes per parameter, 19.2 GB, before activations: it fits on a 32 GB V100. The 8.3B model needs 133 GB: no amount of data parallelism alone will fit it.

In Python:

# 16 bytes per parameter for Adam in mixed precision (the ZeRO paper's count), in GB
16 * 1.2e9 / 1e9  # → 19.2
round(16 * 8.3e9 / 1e9)  # → 133

Why it matters

The two kinds of model parallelism are complementary, as the paper stresses: pipeline across servers, tensor inside a server. The pretraining lesson's section 3 builds all three and shows where each belongs.

3.1 The MLP block · original

Tiny example: two ways to cut

One token, X = [1, 2], and a weight matrix A = [[1, −1], [−1, 1]]. Unsplit: XA = [1 − 2, −1 + 2] = [−1, 1], and GELU gives [−0.159, 0.841].

In words: “multiply the tokens by the first weight matrix, then apply GELU to every number.” (The paper's equation 1.)

Cut by rows. GPU 1 takes A's first row and X's first entry; GPU 2 the second row and second entry.

In words: “cut by rows, each GPU gets a partial sum of every output; the partial sums must be added before GELU, because GELU of a sum is not the sum of GELUs.” (Equation 2 and the sentence after it.)

With the numbers: GPU 1: 1 × [1, −1] = [1, −1]. GPU 2: 2 × [−1, 1] = [−2, 2]. Added first: [−1, 1], GELU → [−0.159, 0.841], correct. GELU first, then added: [0.841 − 0.046, −0.159 + 1.954] = [0.796, 1.796], wrong. So the GPUs would have to talk before the GELU.

Cut by columns. GPU 1 takes A's first column, GPU 2 the second; both need all of X.

In words: “cut by columns, each GPU computes some of the output numbers completely, so it can apply GELU to them on its own.” (Equation 3.)

With the numbers: GPU 1: X · [1, −1]ᵀ = −1, GELU → −0.159. GPU 2: X · [−1, 1]ᵀ = 1, GELU → 0.841. Side by side: [−0.159, 0.841], exactly the unsplit answer, with no communication.

In Python:

import math
def gelu(x):
    return x * 0.5 * (1 + math.erf(x / math.sqrt(2)))
X = [1, 2]
A = [[1, -1], [-1, 1]]
# unsplit: Y = GELU(XA)
XA = [sum(X[i] * A[i][j] for i in range(2)) for j in range(2)]
XA  # → [-1, 1]
[round(gelu(v), 3) for v in XA]  # → [-0.159, 0.841]
# by rows: X_1 A_1 and X_2 A_2 are partial sums of every output
p1 = [X[0] * a for a in A[0]]
p2 = [X[1] * a for a in A[1]]
p1, p2  # → ([1, -1], [-2, 2])
# GELU before adding: wrong
[round(gelu(a) + gelu(b), 3) for a, b in zip(p1, p2)]  # → [0.796, 1.796]
# by columns: each GPU owns whole outputs, GELU runs locally
Y1 = gelu(X[0] * A[0][0] + X[1] * A[1][0])
Y2 = gelu(X[0] * A[0][1] + X[1] * A[1][1])
round(Y1, 3), round(Y2, 3)  # → (-0.159, 0.841)

Reading it: the playground runs the same example on two simulated GPUs. Each card shows what one GPU holds and computes; the last line compares the result with the unsplit Y = GELU(XA) = [−0.159, 0.841]. Cut by rows and each GPU's product is a partial sum of both outputs, so applying GELU locally gives the wrong answer: the GPUs would have to add their partial sums first, a synchronization point. Cut by columns and each GPU finishes one whole output, so GELU runs locally and the answer is exact. That is why Megatron-LM cuts the first matrix by columns.

The second matrix: cut by rows, add once

After GELU, GPU 1 holds Y₁ and GPU 2 holds Y₂: the output is already split by columns. That is exactly the input a row-split of the second matrix B wants. Each GPU multiplies its piece by its rows of B, and one all-reduce adds the two partial results.

In words: “each GPU multiplies its own half of the hidden units by the matching rows of B; adding the two results (one all-reduce) gives the block's output.” (This page's formula for what the paper's Figure 3a draws.)

With the numbers: with B = [[2, 0], [1, 1]]: GPU 1 computes −0.159 × [2, 0] = [−0.317, 0]; GPU 2 computes 0.841 × [1, 1] = [0.841, 0.841]; the all-reduce adds them to [0.524, 0.841], the same as the unsplit YB.

In Python:

import math
def gelu(x):
    return x * 0.5 * (1 + math.erf(x / math.sqrt(2)))
Y1, Y2 = gelu(-1), gelu(1)
B = [[2, 0], [1, 1]]
# each GPU: its hidden unit times its row of B
part1 = [Y1 * b for b in B[0]]
part2 = [Y2 * b for b in B[1]]
[round(v, 3) for v in part1], [round(v, 3) for v in part2]  # → ([-0.317, -0.0], [0.841, 0.841])
# the all-reduce: add the partial results
[round(a + b, 3) for a, b in zip(part1, part2)]  # → [0.524, 0.841]
# unsplit Y B, for comparison
[round(Y1 * B[0][j] + Y2 * B[1][j], 3) for j in range(2)]  # → [0.524, 0.841]

Reading it: data flows left to right. The same X enters both GPUs (the dashed boxes) through f. Each GPU multiplies by its columns of A, applies GELU to the hidden units it owns, and multiplies by its rows of B, all without looking at the other GPU. The two partial outputs meet at g, the block's one all-reduce, before dropout. Read f and g as a pair: in the forward pass f does nothing and g adds; in the backward pass g does nothing and f adds the two GPUs' gradients with respect to X. One all-reduce each way per block.

f and g, in code

The paper's Code 1 defines f as a PyTorch autograd function whose forward returns x unchanged and whose backward all-reduces the gradient; g is the mirror image. They are conjugates: each is the other with forward and backward swapped.

Why it matters today

Column-parallel then row-parallel is the tensor-parallel MLP. The pretraining lesson builds column_parallel_matmul, row_parallel_matmul and tensor_parallel_mlp, and checks that each matches the unsplit answer.

3.2 The attention block · original

Everyday picture

Multi-head attention is already split: each head is an independent small attention with its own query, key and value projections. Hand each GPU a group of whole heads and it can run them start to finish alone.

Tiny example

The 8.3B scaling model has 32 heads on 8 GPUs: 4 heads per GPU. The query, key and value projections are cut by columns so that GPU i's columns are exactly its heads' columns; softmax and the weighted sum stay local. The output projection after attention is cut by rows, and one all-reduce (g) adds the GPUs' partial results, exactly like the MLP.

In Python:

# heads per GPU, and each head's width, for the 8.3B scaling model
heads, gpus, hidden = 32, 8, 3072
heads // gpus, hidden // heads  # → (4, 96)

Reading it: the same skeleton as the MLP. X enters both GPUs through f. The column-split projections give each GPU the queries, keys and values of its own heads; each GPU runs softmax attention for those heads (with attention dropout) entirely locally; the row-split output projection turns each GPU's heads into a partial output; g adds them. The paper's word for what this does to each block is that it “fuses” two matrix multiplies and removes the synchronization point between them.

Why it matters today

Splitting by heads puts a ceiling on this kind of parallelism: you cannot use more GPUs than heads, and the heads must divide evenly. The paper's appendix also finds that more, narrower heads scale a little worse (Appendix D, below). The Attention Is All You Need companion explains the heads themselves.

3.3 Four all-reduces per layer · original

Everyday picture

Put the two blocks together and a whole transformer layer talks four times per training step: once after attention and once after the MLP on the way forward, and once at the start of each block on the way back.

Tiny example

Each all-reduce carries one activation tensor: batch × sequence × hidden numbers. For the 8.3B model with 8 sequences of 1,024 tokens at width 3,072, that is 25.2 million numbers, about 50 MB in 16-bit floats. Four per layer, 72 layers: 288 all-reduces per step, 14.5 GB of messages.

In words: “each layer all-reduces four activation-sized messages per step, two forward and two backward.” (This page's formula for the paper's Figure 4.)

With the numbers: 4 × 8 × 1,024 × 3,072 = 100.7 million numbers per layer; at 2 bytes, 201 MB; times 72 layers, 14.5 GB of messages per step.

In Python:

b, s, H, layers = 8, 1024, 3072, 72
message = b * s * H
round(message / 1e6, 1)  # → 25.2
# C_layer = 4 b s H
C_layer = 4 * message
round(C_layer / 1e6, 1)  # → 100.7
# bytes per step at 2 bytes per number, in GB
round(C_layer * 2 * layers / 1e9, 1)  # → 14.5

Why it matters today

Note what the traffic scales with: the activations, not the weights. Tensor parallelism's cost grows with batch and sequence length, and it happens inside every layer on the critical path, which is why it is kept on the fastest links; the hardware lesson puts numbers on those links.

3.4 The embedding and the loss · original

Everyday picture

The vocabulary table is big (GPT-2's has 50,257 rows) and it is used twice: to look tokens up at the input and, with the same weights (weight tying), to score every token at the output. Megatron-LM splits it by vocabulary: each GPU keeps a slice of the tokens.

Tiny example

At the output, each GPU can score its own slice of the vocabulary. The obvious next step, gathering all the scores (logits) onto every GPU for the cross-entropy loss, would move b × s × v numbers. Instead each GPU computes its slice's part of the loss, and only b × s numbers are exchanged. With b = 8, s = 1,024 and the padded vocabulary of 51,200, that is 419 million numbers against 8,192: 51,200 times less.

In words: “sending losses instead of logits divides the traffic by the size of the vocabulary.” (This page's formula for the paper's §3 argument.)

With the numbers: 8 × 1,024 × 51,200 = 419,430,400 logits against 8 × 1,024 = 8,192 losses; the ratio is v = 51,200.

In Python:

b, s, v = 8, 1024, 51200
V_logits = b * s * v
V_loss = b * s
V_logits, V_loss  # → (419430400, 8192)
V_logits // V_loss  # → 51200

Where does 51,200 come from? For fast matrix multiplies, each GPU's slice of the vocabulary should be a multiple of 128; with up to 8 GPUs, the vocabulary is padded up to a multiple of 128 × 8 = 1,024 (§5.1).

In words: “round the vocabulary up to the next multiple of 128 times the number of GPUs.”

With the numbers: 50,257 / 1,024 = 49.08, rounded up to 50; 50 × 1,024 = 51,200, so 943 padding rows that no real token uses.

In Python:

import math
v, n = 50257, 8
v_pad = math.ceil(v / (128 * n)) * 128 * n
v_pad, v_pad - v  # → (51200, 943)

Why it matters today

Vocabularies have since grown several times larger, so the same trick, never gathering all the logits in one place, matters more now; padding the vocabulary to a friendly multiple is routine.

3.5 Duplicate the cheap parts · original

Everyday picture

Some steps are too cheap to be worth splitting: layer norm, dropout and the residual addition. Rather than have one GPU compute them and send the result, every GPU simply computes them itself on its own copy. Arithmetic is cheap; waiting is not.

Tiny example

A layer norm has 2 × H parameters: 6,144 for H = 3,072, against the roughly 113 million weights of the layer's matrix multiplies (12H², from §5.1 below). Keeping 8 copies of the norm costs almost nothing, and because every GPU holds every value it needs, no updated parameters ever have to be sent: each GPU optimizes its own set.

In Python:

H = 3072
# a layer norm's gain and bias
2 * H  # → 6144
# the layer's matrix-multiply weights, about 12 H², in millions
round(12 * H ** 2 / 1e6)  # → 113

Why it matters

“Recompute or duplicate rather than communicate” recurs throughout large-scale training: activation checkpointing and the FlashAttention backward pass make the same bet.

4.1 Training dataset · original

Everyday picture

Four large sources (Wikipedia, CC-Stories, RealNews and OpenWebText) are pooled; BERT also gets BooksCorpus, which GPT-2 skips because it overlaps with the LAMBADA test. Wikipedia articles that appear in the WikiText103 test set are removed, so the test stays unseen.

Tiny example

Two cleaning steps: drop every document under 128 tokens, and remove near-duplicates with locality-sensitive hashing at a Jaccard similarity above 0.7. Two documents whose combined set of distinct word-shingles (short runs of consecutive words) has 10 members, 7 of them in both, have Jaccard similarity 7/10 = 0.7: not above the threshold, so both stay. With 8 in both, one of them goes. What remains: 174 GB of text.

In Python:

# Jaccard = shared shingles / all distinct shingles
7 / 10, 8 / 10  # → (0.7, 0.8)
# removed only when above 0.7
[j > 0.7 for j in (7 / 10, 8 / 10)]  # → [False, True]

Why it matters today

The same recipe, MinHash and LSH at a Jaccard threshold around 0.7, is built step by step in the pretraining lesson, whose near_duplicate_pairs uses 0.7 as its default.

4.2 Training optimization and hyperparameters · original

Everyday picture

The recipe is standard for its time: mixed precision with dynamic loss scaling, Adam with weight decay 0.01, gradient clipping at a global norm of 1.0, dropout 0.1, and activation checkpointing after every layer. GPT-2 trains on 1,024-token sequences, batch 512, for 300,000 steps, with a peak learning rate of 1.5 × 10⁻⁴ after 3,000 steps of warmup, then cosine decay down to 10⁻⁵. BERT uses batch 1,024 and a learning rate of 10⁻⁴, warmed up over 10,000 steps and decayed linearly over 2 million.

Tiny example: shrinking the last layer of each block

One detail is new. Weights start random with spread 0.02 (the paper writes N(0, 0.02); this page reads 0.02 as the standard deviation), but the weights just before each residual addition are scaled down by 1/√(2N), N being the number of layers. Each layer adds two outputs (attention's and the MLP's) into the residual stream, so 2N contributions pile up. Random contributions add their variances: 144 of them with variance v sum to variance 144v. Scale each by 1/√144 = 1/12 and the sum's variance is back to v.

In words: “the layers that write into the residual stream start with their spread divided by the square root of how many such writes there are.”

With the numbers: the 8.3B model has N = 72 layers, so 2N = 144 residual writes and σ_out = 0.02 / 12 = 0.00167.

In Python:

import math
sigma, N = 0.02, 72
# 2N writes into the residual stream
2 * N  # → 144
sigma_out = sigma / math.sqrt(2 * N)
round(sigma_out, 5)  # → 0.00167
# summed variance of 144 writes, before and after scaling
round(144 * sigma ** 2, 4), round(144 * sigma_out ** 2, 4)  # → (0.0576, 0.0004)

Why it matters today

Scaling residual-branch outputs by depth is a common way to keep very deep stacks stable at the start of training. The deep networks lesson explains why the spread of the signal matters (init_std), and the optimizers lesson builds the schedule (warmup_cosine) and the clipping (clip_by_global_norm).

5.1 Scaling analysis · original

Everyday picture

The test is weak scaling of the model: double the GPUs and double the model, about a billion parameters per GPU, and see whether each GPU stays as busy. Four GPT-2 configurations go from 1.2B on 1 GPU to 8.3B on 8, all with 96 numbers per attention head; then each runs again with 64-way data parallelism, up to 512 GPUs.

Tiny example: checking the parameter counts

A transformer layer holds about 12 × H² weights (4H² in attention, 8H² in the MLP), and the embedding adds v × H. For the largest configuration, 72 layers of H = 3,072 with the padded vocabulary: 12 × 72 × 3,072² + 51,200 × 3,072 = 8.15 + 0.16 = 8.31 billion, the paper's 8.3B. (The rule of thumb is this page's check; the paper gives only the totals.)

In words: “twelve H-squared weights per layer, times the layers, plus one H-wide row per vocabulary entry.”

With the numbers: all four rows of the paper's Table 1: (40 layers, H = 1,536) → 1.21B; (54, 1,920) → 2.49B; (64, 2,304) → 4.19B; (72, 3,072) → 8.31B, against the paper's 1.2, 2.5, 4.2 and 8.3.

In Python:

v = 51200
# (layers, hidden size) from Table 1
configs = [(40, 1536), (54, 1920), (64, 2304), (72, 3072)]
# P ≈ 12 L H² + v H, in billions
[round((12 * L * H ** 2 + v * H) / 1e9, 2) for L, H in configs]  # → [1.21, 2.49, 4.19, 8.31]
# every configuration keeps 96 numbers per head
[H // heads for H, heads in [(1536, 16), (1920, 20), (2304, 24), (3072, 32)]]  # → [96, 96, 96, 96]

The results

The paper gives two figures for the same 512-GPU run: 76% in the abstract and introduction, 74% in §5.1.1. The arithmetic from its own throughput numbers (§1 above) gives 75.6%. Either way, data parallelism on top costs only a few points, because gradients are all-reduced once per step while the model-parallel traffic happens inside every layer.

Why it matters today

Weak scaling of the model, rather than the batch, was the right test for the question people cared about: can a bigger model be trained at nearly the same speed per GPU? The scaling-laws companion explains why people wanted the bigger model in the first place.

5.2 Language modeling results using GPT-2 · original

Everyday picture

Three GPT-2 models, 355M, 2.5B and 8.3B, are trained identically for 300,000 steps. Bigger models learn faster and end lower: the 8.3B reaches a validation perplexity of 9.27.

Tiny example: how long is training?

One pass over the data (an epoch) is 68,507 steps, so 300,000 steps is 4.4 epochs. Each step is 512 sequences of 1,024 tokens, 524,288 tokens, so training sees about 157 billion tokens. For the 8.3B model on 512 GPUs, at 2.10 days per epoch, that is about 9.2 days.

In Python:

steps, per_epoch = 300_000, 68_507
round(steps / per_epoch, 2)  # → 4.38
# tokens per step, and in total (billions)
512 * 1024  # → 524288
round(steps * 512 * 1024 / 1e9)  # → 157
# days for the 8.3B model at 2.10 days per epoch
round(2.10 * steps / per_epoch, 1)  # → 9.2

Reading it: two scores for each model with no fine-tuning (zero-shot), restated from the paper's Table 3. The top group is WikiText103 perplexity, where shorter is better; the bottom group is LAMBADA accuracy, where longer is better. The striped bars are the previous best results. Both scores improve steadily with size, and the 8.3B model passes the previous best on both: perplexity 10.81 against 15.79, accuracy 66.51% against 63.24%. The 2.5B model already beats the old perplexity record.

Checking the test was unseen

The authors count how many 8-word sequences of each test set also appear in the training data: at most 10.8% for WikiText103 (whose own training set already overlaps its test set by 9.09%) and 1.4% for LAMBADA. The paper judges these consistent with earlier work. The same paragraph notes that Turing-NLG, a 17-billion-parameter model, was later trained with Megatron; the ZeRO companion tells that story.

Why it matters today

“Bigger is better, and larger models converge faster” is an early instance of what the scaling laws made quantitative the next year; the GPT-3 companion takes the same curve to 175 billion.

5.3 Bi-directional transformer results using BERT · original

“We further investigated this behaviour and empirically demonstrated that rearranging the order of the layer normalization and the residual connections as shown in Figure 7 is critical to enable the scaling of the BERT-style models beyond BERT-Large.”Shoeybi et al. (2019), §5.3

Everyday picture

Earlier work (ALBERT) had found that BERT got worse beyond 336M parameters. Megatron-LM finds the culprit is plumbing: where the residual connection branches off. If the skip path carries the normalized signal, every layer re-normalizes the running total; if it carries the raw signal, the total flows untouched from bottom to top.

Reading it: both columns have the same boxes in the same order, bottom to top: layer norm, attention, add, layer norm, MLP, add. The only difference is where each thick skip arrow leaves the main line. In (a) it leaves after the layer norm, so what is carried forward and added back is the normalized signal: each block re-normalizes the running total. In (b) it leaves before the layer norm, so the raw input travels up untouched and each block only adds a correction to it. The paper's training curves (Figure 7, right) show (a) training well at 336M parameters but becoming unstable at 752M, while (b) is stable and reaches a lower loss.

Tiny example: the scores

With arrangement (b), bigger is better again. On a 3% held-out set, masked-language-model perplexity falls from 1.58 (336M) to 1.30 (1.3B) to 1.16 (3.9B). On RACE, reading comprehension from exams, the 3.9B model scores 89.5% alone and 90.9% as a 5-model ensemble, against the previous best ensemble's 89.4%.

In Python:

# held-out perplexity for 336M, 1.3B, 3.9B
ppl = [1.58, 1.30, 1.16]
# each step's relative improvement
[round(1 - b / a, 3) for a, b in zip(ppl, ppl[1:])]  # → [0.177, 0.108]
# RACE ensemble: points above the previous best
round(90.9 - 89.4, 1)  # → 1.5

Why it matters today

Arrangement (b) is pre-norm, the layout almost every large transformer now uses. The transformer lesson compares pre-norm and post-norm, and the layer norm companion explains the normalization itself.