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

In data parallelism, the full model is copied to every GPU. Each GPU processes a different mini-batch, computes gradients, and then all GPUs synchronize their gradients via an all-reduce operation. This works beautifully — until the model is too large to fit on a single GPU.

Consider a model with 2 billion parameters in FP16. The weights alone consume 4 GB. Add FP32 optimizer states (Adam stores two moments per parameter), gradients, and activations, and a single GPU needs 40+ GB — exceeding the memory of most GPUs available in 2019.

Pipeline parallelism, used by GPipe, partitions the model into sequential stages placed on different GPUs. But this creates a "pipeline bubble" — GPUs sit idle waiting for activations from previous stages. The deeper the pipeline, the worse the bubble.

What was needed was a way to split each layer across GPUs, so that every GPU participates in every layer simultaneously. This is tensor parallelism — and Megatron-LM made it practical for Transformers.

A Transformer MLP block has two linear layers with a GeLU activation in between. The computation is $Y = \text{GeLU}(XA)$ followed by $Z = \text{Dropout}(YB)$. The weight matrices $A$ and $B$ are the largest memory consumers. The question is: how do you split them across GPUs?

Option 1: split $A$ along its rows. This means each GPU gets a horizontal slice of $A$ and a corresponding vertical slice of $X$. The problem is that GeLU is nonlinear: $\text{GeLU}(X_1 A_1 + X_2 A_2) \neq \text{GeLU}(X_1 A_1) + \text{GeLU}(X_2 A_2)$. You would need to synchronize (all-reduce) before the GeLU — an extra communication step right in the middle of the computation.

Option 2 (Megatron's choice): split $A$ along its columns. Each GPU gets a vertical slice: $A = [A_1, A_2]$. Now each GPU independently computes $\text{GeLU}(X A_i)$, because the GeLU can be applied to each partition separately. No synchronization needed until the very end. Then split $B$ along its rows so it directly accepts the partitioned output from GeLU. The only synchronization is a single all-reduce after the second matrix multiply, right before dropout.

The communication pattern is elegantly simple. Two custom operators handle everything: $f$ is an identity in the forward pass and an all-reduce in the backward pass. $g$ is an all-reduce in the forward pass and an identity in the backward pass. These are conjugates — $f$ and $g$ together ensure that the forward pass produces correct outputs and the backward pass produces correct gradients, with exactly one all-reduce per direction per MLP block. The entire implementation fits in a few lines of PyTorch.

Multi-head attention is inherently parallel. Each attention head operates on a different learned projection of the input — the heads are independent computations that only need to be concatenated at the end. Megatron exploits this by partitioning the Q, K, V projection matrices column-wise so that each GPU computes a subset of attention heads.

Concretely, if a model has 16 attention heads and 4 GPUs, each GPU handles 4 heads. The Q, K, V weight matrices are split so each GPU has its portion of heads. After computing attention, each GPU has partial output. The output linear projection $O$ is split along its rows (just like $B$ in the MLP), so the partial outputs feed directly in without synchronization. A single all-reduce after the output projection combines results — identical to the MLP pattern.

The beauty of this design is that the attention block and the MLP block use the same communication pattern: one $f$ operator at the input, one $g$ operator at the output. The total for a full Transformer layer is just 2 all-reduces in the forward pass and 2 in the backward pass.

The elegance of Megatron's approach lies in two custom autograd operators. Think of them as traffic signals at the boundary of each parallel block:

Operator $f$: In the forward pass, $f$ simply passes the input through unchanged (identity). In the backward pass, $f$ performs an all-reduce to sum gradients across GPUs. This is placed before each parallel block — each GPU receives the full input and computes its own partition.

Operator $g$: In the forward pass, $g$ performs an all-reduce to sum partial outputs across GPUs. In the backward pass, $g$ passes gradients through unchanged (identity). This is placed after each parallel block — the partial results from all GPUs are combined into the full output.

Together, $f$ and $g$ form a conjugate pair. They are the only communication needed. No custom compiler, no special framework — just two custom PyTorch autograd.Function classes.

The authors trained GPT-2 models of increasing size: 1.2B, 2.5B, 4.2B, and 8.3B parameters. All used 8-way tensor model parallelism within each DGX node (8 GPUs connected by NVLink) and data parallelism across nodes.

The 8.3B model on 512 GPUs achieved 15.1 PetaFLOPs sustained throughput — 76% scaling efficiency compared to a strong single-GPU baseline of 39 TeraFLOPs (30% of V100 peak). This was remarkable for 2019: most existing approaches suffered significant efficiency drops at this scale.

A key engineering decision was keeping tensor-parallel groups within a single NVLink domain (one server). NVLink provides 300 GB/s bandwidth between GPUs in a server, while inter-node connections are typically 10–25× slower. By confining the communication-heavy tensor parallelism to the fast interconnect and using the less-communication-intensive data parallelism across nodes, Megatron achieved near-linear scaling.

Scaling BERT beyond 336 million parameters had proven problematic — training became unstable and performance degraded rather than improved. The Megatron-LM authors diagnosed the issue: in the original BERT architecture, layer normalization is applied after the residual connection (post-norm). This means the values flowing through the residual path keep growing with depth, eventually causing instability.

The fix was simple but critical: move layer normalization to before the self-attention and MLP blocks (pre-norm), following the pattern used in GPT-2. With this change, the residual path remains clean — it carries unnormalized values that add linearly — and normalization is applied to the inputs of each sub-layer, keeping activations bounded.

With this rearrangement, the authors successfully trained BERT models up to 3.9 billion parameters — 12× larger than the original BERT-Large — and achieved state-of-the-art results on the RACE reading comprehension benchmark with 90.9% accuracy.

The GPT-2 models showed a clear trend: larger models converge faster and to lower perplexity. The 8.3B model achieved state-of-the-art results across multiple benchmarks:

• WikiText-103: Perplexity of 10.8, down from the previous best of 15.8 — a 32% relative improvement.
• LAMBADA: 66.5% accuracy, up from 63.2% — demonstrating improved long-range language understanding.
• RACE (reading comprehension): 90.9% accuracy with the 3.9B BERT model, surpassing the previous best of 89.4%.

These results confirmed a key hypothesis: given enough compute and an efficient parallelism strategy, simply making Transformer models larger yields significant quality improvements — a theme that would define the next era of NLP with GPT-3 and beyond.

Megatron-LM's impact goes beyond any single model. It established the engineering playbook for distributed LLM training: use tensor parallelism within a node (fast interconnect), pipeline parallelism across a small number of nodes, and data parallelism across the full cluster. Every major training system today — from Meta's LLaMA infrastructure to Google's PaLM training setup — follows this template. The insight that a few all-reduce operations in native PyTorch could replace custom compilers democratized billion-parameter training.