Ulysses Sequence Parallelism: Training with Million-Token Contexts

Summary

Ulysses Sequence Parallelism addresses the memory and compute pressure of training language models on sequences extending beyond tens of thousands of tokens, where attention remains quadratic in sequence length even with FlashAttention. It shards tokens across GPUs, projects local queries, keys, and values, then uses two all-to-all operations to exchange data so each GPU computes a subset of attention heads before restoring sequence-sharded outputs. The post describes integrations with Accelerate, the Transformers Trainer, and TRL's SFTTrainer, including variable-length handling and weighted loss aggregation for uneven valid-token counts. In reported four-GPU tests, SP=4 reduced per-GPU memory sufficiently for 96K-token training on H100 80GB GPUs, reached 13,396 tokens per second at 64K, and matched non-SP canonical token-normalized loss within 0.000005 maximum absolute difference.

Takeaways

1. 1

   Ulysses uses two all-to-all operations per attention layer, with per-GPU communication volume O(n·d/P), while the post reports O(n·d) per GPU for Ring Attention.

2. 2

   On four H100 80GB GPUs, SP=4 supported 96K-token training at 66.0 GB peak memory; the 128K configuration ran out of memory.

3. 3

   Under a matched token budget, canonical token-normalized loss closely matched the non-SP baseline, with SFTTrainer's maximum absolute difference reported as 0.000005.