LIT-022: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Shoeybi et al. (2019) — ARXIV-1909.08053

Key takeaways

• Tensor parallelism as it is now practised. The partition is chosen so the nonlinearity between two GEMMs needs no synchronisation: split the MLP’s first weight matrix column-wise so GeLU applies independently to each piece, then the second row-wise so it consumes that output directly. Same for attention — Q/K/V column-wise, output projection row-wise
• The whole implementation is a conjugate operator pair: f is identity forward and all-reduce backward, g is the reverse. Two all-reduces forward and two backward per transformer layer, in native PyTorch, with no compiler or library changes
• Fuse the parallel logit GEMM into the cross-entropy loss. All-gathering logits communicates b×s×v elements; fusing reduces it to b×s
• 8.3B parameters on 512 GPUs at 15.1 PetaFLOPs, 76% scaling efficiency against a 39-TeraFLOP single-GPU baseline. Vocabulary padded 50,257 → 51,200 so the per-GPU slice is a multiple of 128; hidden size per head fixed at 96 for consistent GEMM shapes
• More attention heads costs model-parallel efficiency — 16 → 32 heads at 8-way parallelism shrinks the per-head GEMMs and grows the softmax. The paper warns explicitly that head count is a systems hyperparameter, not only an architectural one
• Initialisation as stated in the training setup: weights drawn from N(0, 0.02), then the weights immediately before each residual scaled by 1/sqrt(2N), N counting the attention and MLP blocks. Stated as what the runs used, not ablated
• Layer-norm placement decides whether BERT scales. Moving the layer norm relative to the residual (Fig. 7, (a) → (b)) “eliminates instabilities observed using the original BERT architecture” and lowers training loss, which the paper reports as what lets BERT-style models grow past BERT-Large

Standing in the anthology

Read — NOTE-030.

SOTA-060 is sourced to this paper for its initialisation: it is where the N(0, 0.02) draw and the 1/sqrt(2N) residual scaling appear together, as the setup the runs used. The paper does not ablate it.

The reading files no practice for the partition itself. The column-then-row partition is so universal that the record carries it the way it carries backpropagation, which is to say not at all — a gap about the record’s coverage of systems technique rather than an error.

The reading’s live loose end is head count. SOTA-023/SOTA-024 reason about head structure purely in inference memory-bandwidth terms; Megatron says head count also prices in training throughput under tensor parallelism, in the opposite direction, and never quantifies it.

Worth pinning beside SOTA-107. #114 found that practice’s “multiple of 128” was a head dimension and block size in FlashAttention, never a sequence length. Megatron’s 128 is a third thing — the per-GPU vocabulary slice, for logit-layer GEMM efficiency. Three real 128s in three places is exactly how a constant ends up attached to the wrong axis.