Vol. XIV · Parallelism Plate 6 of 11tensor parallelism, Megatron-LM, column parallel, row parallel, activation all-reduce, NVLink domainCE

Splitting Every Layer.

Cut each matrix into column and row blocks, and every device computes a slice of every layer — at the price of an all-reduce inside every layer.

In tensor parallelism the matrices themselves are divided. Shoeybi and colleagues’ Megatron-LM (2019) splits an MLP’s first matrix by columns and its second by rows: each of t devices computes its own slice of the hidden layer independently, multiplies by its rows of the second matrix, and the partial outputs are summed with one all-reduce. Attention is split the same way, by heads. So each layer needs two all-reduces of the activations in the forward pass and two in the backward, each the size of the layer’s output for every token in the micro-batch.

Weights, gradients and optimiser state all shrink by t, and so do most activations. But the communication grows with the tokens, not the parameters, and happens inside every layer where it is hard to hide. The workshop compares one layer’s compute with its all-reduces as t grows, over NVLink inside a server and over the network between servers.

Below: left, one Transformer layer’s forward and backward compute and its four activation all-reduces against the tensor-parallel degree, for one 8 192-token sequence. Right, model states per device.

What you will learn
Workshop
One layer of the selected Llama 3 model on 8 192 tokens: compute 6 × (parameters per layer) × tokens ÷ t at 45% of H100 peak; communication four ring all-reduces of the BF16 activations (2 × tokens × d bytes) among t devices, over NVLink (450 GB/s each way) or a 400 Gb/s network (50 GB/s). Right: model states per device, 16 bytes per parameter ÷ t.

Inside the server, and no further

For one 8 192-token sequence through a Llama 3 70B layer, each all-reduce moves 134.2 MB. Over NVLink with t = 8, the four of them take 20% of the time of the layer’s compute — a cost that can partly be overlapped but not ignored. Over a 400 Gb/s network the same four all-reduces take 161% of it: more communication than computation. Because the traffic scales with tokens × d while compute per device shrinks with t, the ratio only worsens as t grows; at t = 64 over NVLink it is 330%.

That is why tensor parallelism is almost always confined to one NVLink domain: Llama 3 used t = 8, the eight GPUs of one server, which divides the 405-billion-parameter model’s states to 811.7 GB per device — still far too much on its own, so pipeline and data parallelism divide it further. Larger NVLink domains, such as the 72 GPUs of an NVL72 rack, widen the room for tensor and expert parallelism; the principle is the same — the most communication-hungry split goes on the fastest links.

Sequence parallelism (Korthikanti et al., 2023) pairs with tensor parallelism: the parts of a layer that are not split, such as norms and dropout, are split along the sequence instead, turning each all-reduce into a reduce-scatter and an all-gather and dividing their activations by t as well.

Terms

Tensor parallelism
Splitting individual weight matrices across devices.
Column-parallel
Each device holds some columns of a matrix and computes part of its output.
Row-parallel
Each device holds some rows; partial outputs are summed with an all-reduce.
NVLink domain
The set of GPUs joined by NVLink, typically one server of eight.
Tensor parallelism splits every layer and pays two all-reduces per layer each way, traffic that only the fastest links can carry — so it stays inside one server.

Sources

  1. Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism”, arXiv:1909.08053, 2019.
  2. Llama 3 paper, arXiv:2407.21783, section 3.3.2 (tensor parallelism of 8 within a server).