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

Language models get better as they get bigger — until the weights no longer fit on one GPU. Classic data parallelism (DP) replicates the full model on every rank: great when the model is small, useless when even one copy OOMs. Pipeline parallelism slices layers across devices and fights bubbles. Megatron-LM attacks a different axis: intra-layer model parallelism — split the large GEMMs inside each Transformer block across GPUs, keep GeLU and attention heads local, and synchronize with a handful of all-reduces.

The paper’s punchline (abstract): converge Transformers up to 8.3 billion parameters on 512 GPUs, sustain 15.1 PetaFLOPs at 76% scaling efficiency versus a strong single-GPU baseline, and show SOTA-at-the-time WikiText-103 / LAMBADA / RACE numbers for GPT-2- and BERT-style models trained this way.

What this paper explains

Shoeybi et al. give a practical recipe for multi-billion-parameter Transformers in native PyTorch:

1. Name the memory wall. GPT-2 / BERT-style stacks with GeLU and pre-attention LayerNorm still blow single-GPU memory once parameters reach billions.
2. Reject the wrong GEMM split. For $Y=\mathrm{GeLU}(XA)$, splitting $A$ by rows forces a sync before GeLU because $\mathrm{GeLU}(X_1A_1+X_2A_2)\neq\mathrm{GeLU}(X_1A_1)+\mathrm{GeLU}(X_2A_2)$.
3. Use column-then-row parallelism. Split the first MLP (and QKV) GEMM along columns so each GPU applies GeLU locally; split the second GEMM along rows and all-reduce the sum.
4. Mirror that pattern in attention. Column-parallel $Q,K,V$ so heads run locally; row-parallel output projection — no sync between QKV and Softmax.
5. Compose with DP. Model-parallel groups sit inside data-parallel replicas; the approach is orthogonal to pipeline parallelism (complementary, not a replacement).

This is not ZeRO (shard optimizer/grad/param replicas across DP ranks) and not GPipe (stage the depth of the network). It is tensor / model parallelism inside each layer’s matmuls.

Prior limits

• Data parallelism alone keeps communication simple but replicates every parameter, gradient, and optimizer tensor — multi-billion models never leave the station on one device’s memory budget.
• Naive model parallelism that slices activations or weights without respecting nonlinearities inserts extra synchronizations and kills scaling.
• Pipeline parallelism (GPipe-style) partitions layers across stages; useful, but Megatron stresses an intra-layer method that needs no new compiler and works with a few collective ops in PyTorch.
• Cross-node MP pain is real in the broader literature (ZeRO’s related-work notes slow Megatron-style MP across DGX-2 nodes) — Megatron’s contribution is the partitioning pattern and demonstrated scale on large GPU counts, not a claim that NVLink physics changed.

The mechanism

MLP: why column split wins

Write the first feed-forward multiply as

Bad split (rows of $A$). Partition

so $Y=\mathrm{GeLU}(X_1A_1+X_2A_2)$. You must reduce before GeLU.

Good split (columns of $A$). Partition $A=[A_1,A_2]$ and compute

locally. The second GEMM is then row-parallel: it consumes $[Y_1,Y_2]$ without an intervening collective, and an all-reduce ($g$) sums partial outputs before dropout.

Conjugate collectives $f$ / $g$: identity one way, all-reduce the other — a few lines of torch.autograd.Function (paper Code 1).

Self-attention: exploit heads

Column-parallel projections for $Q$, $K$, $V$ place each head’s matmul on one GPU. Softmax attention stays local. The output projection is row-parallel, again avoiding a sync between attention and the projection GEMM.

Layer budget

Fusing pairs of GEMMs this way means a Transformer layer needs only two all-reduces in the forward pass and two in the backward pass (paper Figure 4).

Algorithm / figure walkthrough

1. Form a model-parallel group of size $N_{\mathrm{mp}}$ (and optionally a DP group outside it).
2. Shard MLP weight $A$ by columns across the group; each rank computes local $\mathrm{GeLU}(XA_i)$.
3. Shard the second MLP weight by rows; locally multiply; all-reduce to assemble the MLP output.
4. Shard $W_Q,W_K,W_V$ by columns so each rank owns a subset of heads; run attention locally.
5. Shard the output projection by rows; all-reduce after the projection.
6. Repeat for every layer; train with ordinary Adam/mixed-precision stacks as in the paper’s setup.
7. Scale out: grow $N_{\mathrm{mp}}$ when a layer’s weights no longer fit; grow DP for throughput.

The opening visual is the same arc: ENTER with a full GEMM that OOMs; RUN with column/row shards and local GeLU/heads; LEAVE with a reduced output tensor and a model that fits across the group.

What to notice when reading

• Nonlinearity placement decides the cut. The entire MLP argument is “don’t reduce across a GeLU.” That one inequality drives the column-then-row pattern.
• Attention parallelism is “free” between QKV and Softmax only because heads partition cleanly — head count should play nicely with $N_{\mathrm{mp}}$.
• Orthogonal ≠ free lunch. Megatron + pipeline + ZeRO-style sharding combine in modern stacks; this paper’s claim is the intra-layer primitive, demonstrated up to 8.3B on 512 GPUs.
• LayerNorm placement (BERT path). Beyond parallelism, §5.3 stresses that rearranging LayerNorm in BERT-like models matters for quality as width grows — systems paper with a modeling footnote.
• SOTA numbers are dated. WikiText / LAMBADA / RACE figures are 2019-era comparisons from the abstract; treat them as evidence the method trained competitive models, not as today’s leaderboard.

Results and evidence

Numbers below are from the paper (arXiv:1909.08053) — not re-measured here.

Read the scaling plots in §5.1 for how model-parallel degree interacts with data parallelism; the abstract numbers above are the headline anchors.

Limitations

• Communication fabric matters. Intra-layer all-reduces love high-bandwidth links; weak interconnects erode the 76%-style efficiency the paper reports on its cluster.
• Not a memory panacea alone. Activations, optimizer states, and pipeline bubbles still need ZeRO / checkpointing / PP in today’s largest runs.
• Head / width divisibility. Column-parallel attention assumes you can split heads (and hidden width) across $N_{\mathrm{mp}}$.
• Complementary, not universal. Pipeline and expert parallelism solve different bottlenecks; Megatron does not replace them.
• Benchmark drift. Language-modeling SOTA lines moved after 2019 — use the paper’s tables as historical evidence of successful large-model training.

How to read the paper

1. Abstract + §3 (model-parallel Transformers) — the GEMM partitioning argument and Figure 3.
2. Figure 4 — count the collectives per layer.
3. §5.1 — scaling with hybrid model + data parallelism.
4. §5.2–5.3 — GPT-2 / BERT quality results and the LayerNorm placement note.
5. Appendix B — hybrid MP+DP and RNG details if you implement.

Knowledge check

Keep reading / Sources

• Paper: arXiv:1909.08053 — Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Related Daily: ZeRO (optimizer/state sharding), GPipe (pipeline micro-batches), Ring Attention (sequence-parallel KV)
• LayerNorm background: Ba, Kiros, Hinton — Layer Normalization