Sequence and Context Parallelism

Every other parallelism strategy splits the model or the batch. Neither helps with the axis that dominates long-context training. Activations scale with sequence length, and attention's intermediate state scales worse than linearly, so a model that trains comfortably at 8k tokens can be impossible to train at 128k on the same hardware with the same batch size and the same parallel layout. Sequence parallelism splits the sequence itself.

The memory that data and tensor parallelism cannot touch

Data parallelism replicates the model and shards the batch, so per-device activation memory is unchanged for a fixed local sequence. Tensor parallelism shards weight matrices, which cuts parameter and some activation memory but leaves per-token activations proportional to sequence length. With a batch of one long document, data parallelism has nothing left to split.

The problem is concentrated in attention. The scores matrix for one head is \(S \times S\) for sequence length \(S\), so going from 8k to 128k multiplies that term by 256. FlashAttention removes the need to materialise it in HBM by tiling the computation in SRAM, which is why 128k training is possible at all, but the KV tensors and layer activations still scale linearly in \(S\) and still have to live somewhere.

Two ways to split a sequence

Ring Attention shards the sequence across devices and passes KV blocks around a ring while each device computes attention over its local queries (Liu, Zaharia & Abbeel, 2023, Ring Attention with Blockwise Transformers, arXiv:2310.01889). Device \(i\) holds queries for its shard permanently, and over \(N\) steps receives every other device's keys and values in turn, accumulating attention output with the same online-softmax rescaling FlashAttention uses. The communication is peer-to-peer and can overlap with compute, so with enough blockwise work per step the ring transfer hides entirely. Context length then scales with device count, which is where "near-infinite context" in the title comes from.

DeepSpeed-Ulysses takes the other route: shard by sequence outside attention, then use all-to-all collectives to switch to a head-sharded layout for the attention computation itself, and switch back afterwards (Jacobs et al., 2023, arXiv:2309.14509). Each device then computes full-sequence attention for a subset of heads. Communication volume is lower than the ring for many shapes, but the degree of parallelism is capped by the head count, which grouped-query attention has been busy reducing.

The two compose, and production stacks increasingly run them together on different axes of the device mesh.

What it does not change

Sequence parallelism is mathematically transparent. The loss, the gradients and the resulting weights are identical to single-device training up to floating-point non-determinism. This is what distinguishes it from sliding-window attention or other context-length tricks: nothing is approximated, only redistributed. A model trained with context parallelism has no architectural marker of it.

Note that "sequence parallelism" also names a narrower Megatron-LM technique that shards LayerNorm and dropout activations along the sequence axis to complement tensor parallelism. Both meanings are current, and papers usually rely on context to disambiguate. The long-context variant is now more often called context parallelism.

When it breaks

• Causal masking makes the ring load-imbalanced. With a causal mask, the device holding the last sequence shard attends over everything before it while the first device attends over almost nothing. Naive sharding leaves early devices idle for most of each step. Production implementations interleave or zigzag the assignment so each device gets a mix of early and late positions.
• Overlap depends on block size. The ring hides its transfer only if each step's compute exceeds the KV block transfer time. Small blocks, small models, or slow interconnects break the overlap and the ring becomes a latency chain of \(N\) hops.
• Ulysses is capped by head count. With 8 KV heads, an all-to-all attention layout parallelises across at most 8 devices, which is a hard ceiling regardless of cluster size.
• It composes with, and constrains, everything else. Context parallelism consumes a dimension of the device mesh that tensor or pipeline parallelism wanted. See composing parallelism strategies for how the mesh gets allocated.