For very long sequences, activation memory grows with sequence length and eventually blows up even after other parallelism. Context parallelism splits the sequence dimension across GPUs, so each device holds only a slice of the tokens. The hard part is attention, where every token must see every other token.
The problem
Most of a Transformer is pointwise along the sequence, so splitting tokens across GPUs is trivial for the MLP and norms. Attention is the exception: a query on GPU 0 needs keys and values from tokens held on every other GPU.
Ring Attention
- Arrange GPUs in a ring, each holding one block of queries, keys, and values.
- Compute attention against the local KV block, then rotate the KV blocks around the ring.
- After one full loop, every query has attended to every key, and no GPU ever held the full sequence.
- KV communication overlaps with attention compute, so the cost is largely hidden.
Load imbalance from causal masking
With a causal mask, early tokens attend to few keys and late tokens attend to many, so a naive block split leaves some GPUs idle. Zig-zag (striped) ring attention assigns each GPU a mix of early and late tokens so every device does roughly equal work.
When to use it
- Reach for context parallelism only when the sequence itself is the memory bottleneck (long-context training or fine-tuning).
- It composes with data, tensor, and pipeline parallelism as one more axis in 3D Parallelism.
- Related idea at inference time: see KV Cache and long-context serving.