Vol. XIV · Parallelism Plate 8 of 11context parallelism, ring attention, sequence splitting, long context, key–value exchange, load balancing, causal maskCE

Splitting the Sequence.

When one sequence is too long for one device, give each device a piece of it — and pass the keys and values around.

At 131 072 tokens, the activations of one sequence through a large model run to tens of gigabytes per device even after tensor and pipeline parallelism. Context parallelism splits the sequence itself across c devices. Everything in a layer except attention works token by token, so it splits for free. Attention needs every earlier key and value: in ring attention (Liu et al., 2023) each device keeps its block of queries and passes its block of keys and values to the next device in a ring, computing attention against each block as it arrives and combining the results with the online softmax (plate XVI.3). After c − 1 passes every query has seen every key.

Two details make it work. With grouped-query attention the keys and values are small — 8 heads against 128 query heads in Llama 3 405B — so passing them costs little next to the attention compute and can be overlapped. And the causal mask makes a naive split unbalanced: the device holding the end of the sequence attends to everything, the one holding the start to almost nothing. Splitting the sequence into 2c chunks and giving device k chunks k and 2c − 1 − k equalises the work. The workshop measures all three effects.

Below: left, activation memory per device for one sequence of the chosen length through Llama 3 405B (tensor parallelism 8, pipeline 16, FlashAttention) against the context-parallel degree. Right, causal attention work per device for a contiguous split and a zig-zag split.

What you will learn

Explain why long sequences need context parallelism.
Describe ring attention.
Compare key–value traffic with attention compute.
Balance causal work with a zig-zag split.

Workshop

Llama 3 405B with tensor parallelism 8 and pipeline parallelism 16 (about 8 layers per stage), FlashAttention, one sequence per micro-batch; activation memory per device from the Korthikanti et al. estimate divided by t and c. Ring traffic per layer: this device’s share of the BF16 keys and values (1 of 8 KV heads under t = 8), (c − 1)/c of it sent over a 400 Gb/s network; attention compute per layer per device, causal, at 45% of H100 peak.

One sequence, sixteen devices

A single 131 072-token sequence through Llama 3 405B, even with tensor parallelism 8 and pipeline parallelism 16, leaves 71.9 GB of activations per device — almost an H100 by itself, before the model states and before the pipeline’s in-flight micro-batches. With c = 16 it is 4.5 GB. Passing each device’s keys and values round the ring takes 12.7% as long as the attention computed on them, because grouped-query attention makes them small, so the transfer hides behind compute. At 8 192 tokens none of this is needed: one sequence takes 4.5 GB.

The contiguous split makes the busiest device do 1.94 times the average attention work at c = 16, and every device waits for it; the zig-zag split makes the work equal. Llama 3 used context parallelism of 16 for its long-context stage, with a variant that all-gathers the keys and values instead of passing them round a ring — simpler to combine with document masks, and affordable for the same reason: under grouped-query attention, keys and values are a small fraction of the activations.

Context parallelism and the sequence parallelism of plate XIV.6 are easily confused: the latter splits the non-attention parts of a layer alongside tensor parallelism; the former splits attention itself, and is needed only when sequences are long.

Terms

Context parallelism: Splitting one sequence across devices, including its attention.
Ring attention: Passing key–value blocks round a ring of devices so every query sees every key.
Zig-zag split: Giving device k chunks k and 2c − 1 − k to balance causal work.
Load imbalance: Uneven work across devices, so all wait for the busiest.

Context parallelism splits a single long sequence across devices; small grouped-query keys and values make the ring exchange cheap, and a zig-zag split balances the causal mask.

Sources

Liu, Zaharia and Abbeel, “Ring Attention with Blockwise Transformers for Near-Infinite Context”, arXiv:2310.01889, 2023.
Llama 3 paper, arXiv:2407.21783, section 3.3.2 (context parallelism of 16 for 131 072-token sequences; all-gather of keys and values).