Megatron-LM

The Problem

Even with memory optimizers like ZeRO, some models (like GPT-3 at 175B parameters) are so astronomically large that even a single layer's parameters and activations cannot fit into the memory of a single GPU.

Furthermore, doing all the math for a massive matrix multiplication on one GPU is slow. Researchers needed a way to split the model itself across GPUs. Previous attempts at Model Parallelism usually involved "Pipeline Parallelism"—putting Layer 1 on GPU A, and Layer 2 on GPU B. But this causes severe "bubble" inefficiencies where GPU B sits completely idle waiting for GPU A to finish.

The Idea

The NVIDIA researchers developed Tensor Parallelism (TP) specifically tailored for the Transformer architecture. Instead of putting different layers on different GPUs, they put half of the same layer on GPU A and the other half on GPU B, allowing both GPUs to work on the exact same token simultaneously.

They realized that the Transformer architecture consists primarily of two components: Self-Attention and MLPs. Both of these are essentially just massive matrix multiplications. By clever linear algebra, they could slice the matrices so that the GPUs could do their half of the math independently, requiring only a single synchronization step at the end.

How It Works

Consider the two linear layers in an MLP: $Y = \text{GeLU}(X A) B$.

1. Column Parallelism: They slice matrix $A$ vertically. GPU 1 gets $A_1$, GPU 2 gets $A_2$. Both GPUs receive the input $X$. GPU 1 computes $Y_1 = \text{GeLU}(X A_1)$ and GPU 2 computes $Y_2 = \text{GeLU}(X A_2)$. Crucially, they do not need to talk to each other to do this.
2. Row Parallelism: They slice matrix $B$ horizontally ($B_1$ and $B_2$). GPU 1 multiplies its intermediate result $Y_1$ by $B_1$. GPU 2 multiplies $Y_2$ by $B_2$.
3. Synchronization (All-Reduce): At the very end, the GPUs simply sum their final outputs together ($Y_1 B_1 + Y_2 B_2$).

This exact same logic (column parallel followed by row parallel) is applied to the Multi-Head Attention blocks, splitting the attention heads across GPUs.

Why It Mattered

Megatron-LM proved that intra-layer Model Parallelism was highly efficient. By keeping the communication synchronized and minimal (only one All-Reduce per block), they achieved massive scaling efficiency. It allowed NVIDIA to train some of the largest models in the world with minimal overhead, utilizing the ultra-fast NVLink connections between GPUs in a single server node.

What Came After

Megatron-LM's Tensor Parallelism is an absolute requirement for training or serving models over ~70B parameters today. It is almost always combined with Data Parallelism (ZeRO) and Pipeline Parallelism to create 3D Parallelism, which is the gold-standard architecture used to train GPT-4, Llama 3, and every other frontier model.