One dimension of a training step is still unsplit, the sequence. Train Llama 3 70B on 128K-token documents, as the last stage of long-context training does, and one sequence's activations, even with FlashAttention, tensor parallelism, and sequence parallelism over 8 GPUs, come to $34Td/t$ bytes per layer for $T = 131{,}072$ tokens, width $d = 8192$, and tensor degree $t = 8$, about 4.6 GB, times 80 layers, 365 GB on every GPU. The other strategies cannot help. Data parallelism needs at least one whole sequence per GPU, tensor parallelism is capped at the node, and pipelining splits layers, not tokens (under 1F1B the first stage holds all $L$ layers' worth of activations anyway, §10.5).

Context parallelism splits the sequence. It is easy for everything except attention, where every token needs every earlier token. This section is about the systems side of that split, where each piece of the sequence lives, what has to travel between GPUs, what it costs, and when it hides. The attention algorithm that keeps the split exact, ring attention built on the online softmax, belongs to the chapter on efficient attention (§12.2 and §12.7), and we recap only the part we need. We finish by placing the experts of a mixture-of-experts model and by putting every strategy of the chapter on one grid of GPUs.

Splitting the sequence, and the one layer that objects

Cut a sequence of $T$ tokens into $n_{\text{cp}}$ contiguous chunks and give one chunk to each GPU of a context-parallel group. Embeddings, norms, the $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ and output projections, the feed-forward network, and the loss all act on each token independently, so each GPU runs them on its own $T/n_{\text{cp}}$ tokens with no communication, and every activation of those layers is split $n_{\text{cp}}$ ways. The weights are not split, so context parallelism is usually combined with FSDP over the same GPUs, which shards them.

Attention is the exception. The query of token $i$ must see the keys and values of every earlier token, most of which live on other GPUs. There are two ways to get them there.

Gather them. All-gather the keys and values so that every GPU has the full $\mathbf{K}$ and $\mathbf{V}$, then compute attention for the local queries. Llama 3 does this. Its reasoning is that with grouped-query attention the keys and values are much smaller than the queries (8 key-value heads against 64 query heads for the 70B model), that the all-gather is cheap next to the quadratic attention computation, and that gathering makes arbitrary masks, such as document masks for packed sequences, easy to support. The backward pass is the mirror image. Each GPU computes gradients for the full keys and values from its own queries, and a reduce-scatter sums those contributions and leaves each GPU the gradient of its own chunk, the same conjugate pair of collectives as sequence parallelism (§10.4). The costs are that the communication sits on the critical path and that each GPU briefly holds the full $\mathbf{K}$ and $\mathbf{V}$.

Pass them around a ring. Ring attention (Liu et al. 2023) never gathers. The GPUs form a ring, and each GPU starts with its own key-value block. In each of $n_{\text{cp}}$ steps, every GPU computes the attention of its queries against the block it currently holds, and at the same time sends that block to its right neighbor and receives the next one from its left. After $n_{\text{cp}}$ steps every query has seen every block, and no GPU ever held more than two key-value blocks at once (the one it computes with and the one arriving).

How can attention computed one block at a time equal attention over the whole sequence? Each GPU keeps, for every query, a running maximum of the scores seen so far, a running sum of their exponentials, and a running weighted sum of the values, and it rescales the old totals whenever a block brings a larger score. This is the online softmax, and the result is exact, not an approximation. §12.2 derives it, and §12.7 builds ring attention from it and tests it against ordinary attention. For placing the work, all we need is the communication pattern. Here it is for four GPUs.

What context parallelism communicates

The cell already shows the surprise. Count the forward pass of one attention layer on one GPU, with $T$ tokens split over $n_{\text{cp}}$ GPUs, so that each holds $n = T/n_{\text{cp}}$ of them, and with key-value width $d_{kv} = n_{kv}d_k$ for $n_{kv}$ key-value heads of width $d_k$ (equal to $d$ without grouped-query attention). One GPU's keys and values are $2n\,d_{kv}$ numbers, $4n\,d_{kv}$ bytes in 16 bits. A ring all-gather (§10.2) makes each GPU receive the other $n_{\text{cp}} - 1$ chunks, and ring attention makes each GPU receive $n_{\text{cp}} - 1$ blocks, one per step after the first. Both come to

In plain words, both methods deliver every other GPU's keys and values to every GPU once, so they move the same bytes, almost the full sequence's keys and values, however many GPUs share the sequence.

So the choice between them is not about volume. It is about timing and memory. The all-gather moves everything before the attention kernel runs, holds the full keys and values, and works with any mask and one ordinary attention kernel. The ring spreads the same bytes over $n_{\text{cp}}$ steps, overlaps each transfer with a step of attention, and holds only two blocks. In the backward pass both roughly double the traffic, because the key and value gradients must travel back to their owners, as a reduce-scatter in the all-gather version and as a second circulating block in the ring.

Is the exposed all-gather affordable? Take Llama 3 405B's long-context stage ($d = 16{,}384$, 128 query heads and 8 key-value heads of width 128, 126 layers), with $T = 131{,}072$, $n_{\text{cp}} = 16$, and $t = 8$. Tensor parallelism has already split the heads, so each GPU holds one key-value head, and the context-parallel group spans 16 nodes, so the all-gather runs over InfiniBand. The cell compares it with the GPU's forward compute for the same layer, causal attention for its 8,192 queries plus its share of the weight products, at 400 TFLOP/s.

With grouped-query attention the all-gather moves 63 MB per GPU per layer, about 1.3 ms, under 5% of the 27 ms of compute, so leaving it exposed costs little, which is exactly Llama 3's argument. Without grouped-query attention the same layer would move 16 times as much, 20 ms, about three quarters of the compute, and then overlapping it with a ring would matter.

When the ring hides its communication

The ring overlaps each transfer with a step of attention, so the question is whether the step is long enough. With blocks of $n$ tokens, the attention of a fully visible block costs $4n^2 d$ FLOPs ($2n^2d$ for $\mathbf{Q}\mathbf{K}^\top$ and $2n^2d$ for the weighted values, over all heads), and the key and value blocks are $4n\,d_{kv}$ bytes. With a throughput of $F$ FLOP/s per GPU and $\mathcal{B}$ bytes per second on the link to the next GPU, the send hides behind the compute when

In plain words, each GPU's chunk of the sequence must be at least the GPU's FLOPs-to-bandwidth ratio in tokens, and grouped-query attention lowers the bar by the ratio of query heads to key-value heads.

Liu et al. state the condition as a block size of at least $F/\mathcal{B}$ (they assume $d_{kv} = d$), and §12.7 derives it per head. For H100s at peak throughput that is about 2,200 tokens per GPU over NVLink and about 20,000 over InfiniBand, and eight times less with Llama-style grouped-query attention. A 128K-token sequence over 16 GPUs gives 8,192 tokens per GPU, comfortably above the NVLink bar and, with grouped-query attention, above the InfiniBand bar too. The backward pass circulates the key and value gradients as well, about twice the bytes, but it also does about twice the attention FLOPs, so the condition barely changes.

The causal mask unbalances the sequence

Whichever method delivers the keys and values, a causal mask makes contiguous chunks unfair. GPU $i$'s queries need only the chunks at or before its own, so it computes $i + 1$ blocks of attention while GPU 0 computes one, and that one is only half visible. Counting a diagonal block as half, the last GPU does $2n_{\text{cp}} - 1$ times the work of the first, and every GPU waits for the slowest. The fix is a placement decision. Give every GPU a mix of early and late tokens. Split the sequence into $2n_{\text{cp}}$ chunks and give GPU $i$ chunks $i$ and $2n_{\text{cp}} - 1 - i$, one from the front and its mirror from the back. This zigzag layout is what Llama 3 uses. GPU 0 gets the first chunk (cheap, it sees almost nothing) and the last (expensive, it sees everything), and every GPU's total comes out equal. Striped attention (Brandon et al. 2023), which §12.7 simulates, achieves the same by assigning tokens round-robin. The zigzag keeps each GPU's tokens in two contiguous runs, which suits kernels and rotary position ids. The widget shows the work of every GPU at every ring step.

def chunk_assignment(c, layout):
    return [[i] for i in range(c)]

def causal_work(assignment):
    return [float(len(chunks)) for chunks in assignment]

Ulysses, splitting by heads instead

DeepSpeed-Ulysses (Jacobs et al. 2023) takes another route to the same goal. Attention heads are independent, so if each GPU held all tokens but only some heads, it could compute exact attention locally with an ordinary kernel. Ulysses switches layouts with an all-to-all (§10.2). Each GPU starts with $T/n_{\text{cp}}$ tokens of $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ for all heads, and one all-to-all trades them so that each GPU has all $T$ tokens for $h/n_{\text{cp}}$ heads. After local attention, a second all-to-all trades the outputs back to the token split.

Count its bytes the same way. An all-to-all keeps $1/n_{\text{cp}}$ of each GPU's buffer and sends the rest. Each GPU's share is $n = T/n_{\text{cp}}$ tokens of $\mathbf{Q}$ (width $d$), $\mathbf{K}$ and $\mathbf{V}$ (width $d_{kv}$ each), and the output (width $d$), so per attention layer in the forward pass

In plain words, Ulysses moves only each GPU's own share of the activations, which shrinks as GPUs are added, while the ring and the all-gather deliver nearly the whole sequence's keys and values to every GPU, so Ulysses wins once the group is larger than $1 + h/n_{kv}$ GPUs.

Without grouped-query attention ($h = n_{kv}$) that is any group larger than 2. With 8 query heads per key-value head it is a group larger than 9, and here the head count bites. Ulysses needs each GPU to own whole heads, so $n_{\text{cp}}$ cannot exceed the number of heads, and the keys and values can be split only $n_{kv}$ ways (beyond that they are replicated, which costs more). For a Llama model with 8 key-value heads, Ulysses tops out at 8 GPUs, just short of the point where it would beat the ring. That, together with the all-to-all's sensitivity to topology and its place on the critical path, is why grouped-query models favor the all-gather or the ring, while models with many key-value heads often run Ulysses within a node and the ring across nodes.

def cp_comm_bytes(T, d, d_kv, n_cp, method):
    return 4 * T * d_kv            # placeholder: ignores the method and the group size

def ulysses_breakeven(h, n_kv):
    raise NotImplementedError

Where the experts go

A mixture-of-experts model (Chapter 13) adds one more thing to place. DeepSeek-V3 has 256 routed experts in each of its MoE layers, far too many to replicate on every GPU, so they are split, $E/n_{\text{ep}}$ of the $E$ experts on each GPU of an expert-parallel group of $n_{\text{ep}}$ GPUs, while attention and the other shared layers stay data parallel. Tokens travel to their experts and back with two all-to-alls per MoE layer in the forward pass and two in the backward pass.

For the layout, three facts matter. Expert parallelism does not need GPUs of its own. It reuses GPUs that are data-parallel (and context-parallel) for the rest of the model, so the expert-parallel group is carved out of those axes, and each expert's gradient is reduced only across the GPUs that hold a copy of that expert. All-to-all is the collective most sensitive to topology, because every GPU talks to every other at once, so the group is kept inside a node or a few nodes. And expert parallelism competes with tensor parallelism for the job of splitting the feed-forward weights. DeepSeek-V3 used 64-way expert parallelism over 8 nodes and no tensor parallelism at all, while models with a few large experts often split each expert with tensor parallelism instead. How dispatch works, what it costs per token, and how DeepSeek-V3 caps each token's traffic are the subject of §13.4.

Putting the dimensions on one grid

We now have five ways to split a training step, data (with ZeRO or FSDP), tensor (with sequence), pipeline, context, and expert parallelism. A real run uses several at once by arranging its GPUs as a multidimensional grid, a device mesh, with one axis per strategy. Every GPU has a coordinate on each axis and belongs to one group per axis, the GPUs that differ from it only along that axis. The product of the degrees is the number of GPUs, for example $n_{\text{gpu}} = t \cdot n_{\text{cp}} \cdot p \cdot n_{\text{dp}}$, with experts folded into the data and context axes.

The order of the axes decides which groups share a node. Llama 3 orders them, from innermost (adjacent ranks) to outermost, as tensor, context, pipeline, data. Tensor parallelism, the most bandwidth-hungry, gets consecutive ranks inside a node. Data parallelism, whose traffic is the most latency-tolerant and overlappable, spans the whole cluster. Concretely, rank $r$'s coordinates are its digits in a mixed-radix number system whose fastest digit is the tensor coordinate.

Axes can also be combined. The $n_{\text{cp}}$ GPUs of a context-parallel group hold the same weights and, once their gradients are summed, the same gradient, so the gradient reduction and the ZeRO or FSDP sharding run over the data and context axes together, $n_{\text{dp}}n_{\text{cp}}$ GPUs, as Megatron's distributed optimizer and PyTorch's torchtitan do. §10.7 uses this to keep a long-context run inside memory.

def parallel_groups(tp, cp, pp, dp):
    n = tp * cp * pp * dp
    return {"tp": [list(range(n))], "cp": [], "pp": [], "dp": []}

In PyTorch

Both pieces need several processes and are shown for reference. PyTorch builds the mesh in one call, outermost axis first, and the last axis gets consecutive ranks. Slicing it by name gives the group of each strategy, and flattening two axes gives the combined group of the previous paragraph.

Context parallelism for attention written with F.scaled_dot_product_attention is a context manager. It shards the listed buffers along their sequence dimension in place (with a load-balanced, zigzag-style split for causal attention) and swaps the attention call for a context-parallel one over cp_mesh. Every tensor indexed by position must be in the list, or the rotary embeddings will use local positions. The default delivers the keys and values with an all-gather, and set_rotate_method("alltoall") passes them around the ring instead.

Read more

• Liu, Zaharia, and Abbeel (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. The ring, the overlap condition, and the memory analysis.
• Jacobs et al. (2023). DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. Sequence-to-head all-to-alls and their communication volume.
• Brandon et al. (2023). Striped Attention: Faster Ring Attention for Causal Transformers. Balancing the causal mask by striping tokens across devices.
• Llama Team, AI at Meta (2024). The Llama 3 Herd of Models. Section 3.3.2 describes 4D parallelism, all-gather context parallelism, and the mesh order.
• DeepSeek-AI (2024). DeepSeek-V3 Technical Report. Section 3.2 gives a frontier MoE layout with expert parallelism and no tensor parallelism.
• PyTorch. Introduction to Context Parallel. The context_parallel API and its all-gather and all-to-all rotation methods.
• Tazi et al. (2025). The Ultra-Scale Playbook. Its context parallelism chapter, with diagrams of both ring and zigzag layouts.

Watch next

▶

▶