Training a Model with 16,384 GPUs: Why Is Data Parallelism Far from Enough?

Training frontier AI models requires thousands of GPUs to work together. Meta used 16,384 GPUs to train Llama 3 405B. The training workload must be split along different axes and then synchronized and reassembled.

Jia-Bin Huang made a 40-minute video explaining 5D parallelism: data, tensor, context, pipeline, and expert. These five types of parallelism split the workload along the five directions of data, width, sequence length, layers, and experts.

The video starts with data parallelism and first reviews the three stages of ZeRO: sharding away redundancy in optimizer states, gradients, and weights. But data parallelism has a hard limitation: it cannot split a single sequence. Llama 3 uses a global batch of 2,048 sequences, each with 8,192 tokens. Spread across 16,384 GPUs, this means every 8 GPUs together process one sequence. Data parallelism stops here; it cannot split further.

A bigger problem is activations. For a sequence of 131,072 tokens, a single layer's input in the 405B model requires 4.3GB, and over 126 layers this accumulates to 541GB. This does not yet include the extra memory needed for backpropagation. Add communication overhead, and when the number of GPUs doubles, ZeRO-3 can actually become slower.

The training task can be viewed as a cube with three dimensions: tokens, sequences, and layers (tokens × sequences × layers). Each small cell is the activation value of a token at a certain layer. Data parallelism only splits the batch dimension, so it is necessary to split along other axes.

Tensor parallelism splits along width. Linear layers have two ways to split: column-linear, using all-gather to concatenate the outputs; and row-linear, using all-reduce to sum partial results. In an MLP layer, the first matrix is split by columns and the second by rows, so the entire layer needs only one all-reduce. The attention layer works similarly: the QKV projection is split by columns, and the output projection is split by rows. Note that Llama 3 405B has 128 query heads but only 8 KV heads, so the maximum tensor parallel size is 8.

Context parallelism addresses ultra-long sequences. When a sequence is too long for a single server, ring attention sends KV blocks around a ring and, together with online softmax, enables long-sequence training across servers.

Pipeline parallelism splits along layers. The naive schedule has large bubbles; GPipe introduces micro-batches to fill the bubbles; 1F1B interleaves forward and backward passes; Zero-bubble and DualPipe reduce bubbles even further.

Expert parallelism is for MoE (mixture-of-experts) models. Different experts are placed on different GPUs, and tokens are sent to the corresponding experts via all-to-all.

At the end, the video breaks down the 16,384-GPU layout for Llama 3 and also provides a three-step quick reference table to teach you how to design your own parallelism strategy. One viewer said this is the only channel they have to watch at 0.5x speed. The diagrams are dense.

The full video is here:

Original tweet: https://x.com/jbhuang0604/status/2114680942619226176