Vol. XIV · Parallelism Plate 10 of 113D parallelism, 4D parallelism, layout, tensor, pipeline, data parallelism, MFU estimate, Llama 3 405BCE

Putting It Together.

Tensor inside the server, pipeline across servers, data around the outside — and every choice costs memory, bubble or traffic somewhere.

No single scheme trains a 405-billion-parameter model. Llama 3 used four at once on 16 384 H100s: tensor parallelism of 8 over each server’s NVLink, pipeline parallelism of 16 across servers, context parallelism for its long-context stage, and data parallelism with sharded optimiser state over the rest. The product of the degrees is the number of GPUs, so raising one lowers another, and each trades differently: tensor parallelism costs bandwidth, pipelines cost bubble, data parallelism costs memory unless sharded, and the global batch must be shared among the data-parallel replicas as micro-batches for the pipeline.

The workshop lets you choose a layout for a Llama 3 model on a given number of GPUs with a global batch of 16 million tokens (2 048 sequences of 8 192). It reports memory per device, the bubble, tensor- and data-parallel communication, and an estimate of the resulting model FLOPs utilisation, starting from an assumed 60% kernel efficiency and multiplying in each loss.

Below: left, memory per device for the chosen layout by part, against 80 GB. Right, the estimated utilisation, step by step: kernel efficiency, then after the pipeline bubble, tensor-parallel communication, data-parallel communication and recomputation.

What you will learn
Workshop
Memory per device from the activation and state accounting of plates XIV.1–XIV.8 (FlashAttention, sequence parallelism with t, one sequence per micro-batch, min(m, p) micro-batches in flight). Utilisation: 60% kernel efficiency (an assumption); bubble with v interleaved chunks per device; four activation all-reduces per layer over NVLink for t ≤ 8 or the 400 Gb/s network beyond, not overlapped; the gradient all-reduce over the network, overlapped with the step; r = 6/8 if every layer is recomputed. A rough model, not a simulator.

Llama 3’s layout, roughly

The defaults are Llama 3 405B’s published layout on 16 384 GPUs: t = 8, p = 16 and so 128 data-parallel replicas, each with 16 sequences per step as micro-batches. With 4 interleaved chunks the bubble is 19.0%, the tensor-parallel all-reduces cost a little more, and the estimate comes to 43% — near the 41% Meta reported for this configuration, which is a check on the model, not proof of it. Memory is the tight constraint. Llama 3’s FSDP sharded the optimiser state and gradients (ZeRO-2 here), which leaves 78.6 GB per device, just under 80 GB, almost all of it activations; Meta reports training 8K sequences without activation checkpointing after profiling memory and freeing every tensor it could. Recomputing every layer would bring memory down to 10.9 GB but the estimate down to 32%.

Each alternative fails somewhere. Without interleaving the bubble rises to 48.4% and utilisation to 27%. Stretching tensor parallelism to 16 across the network, with 8 pipeline stages, drops it to 16%. Without a pipeline each device needs 173.7 GB; without any sharding, 122.6 GB. A 70-billion-parameter model on 2 048 GPUs fits with t = 8 and 4 stages at 27.7 GB and an estimated 43%. Choosing a layout is a search over these trade-offs, and real teams run it with profilers, not formulas.

The order of the dimensions matters as much as their sizes: the innermost (tensor) must sit on the fastest links, the next (context or pipeline) on the next fastest, and data parallelism, which communicates least often, on the slowest.

Terms

3D parallelism
Tensor, pipeline and data parallelism combined.
Layout
The degrees of each parallelism and how they map onto the hardware.
Interleaving
Giving each pipeline device several non-adjacent chunks of layers.
Utilisation estimate
Kernel efficiency multiplied by the losses from bubble and communication.
A layout multiplies tensor, pipeline, context and data parallelism to the GPU count, and each degree trades memory, bubble or traffic; Llama 3’s choice sits where all three are tolerable.

Sources

  1. Llama 3 paper, arXiv:2407.21783, section 3.3.2 and Table 4 (TP 8, PP 16, CP 1 or 16, DP 64–128; BF16 MFU 38–43%).
  2. Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM”, SC 2021.