Ulysses Sequence Parallelism: Training with Million-Token Contexts

This post explains Ulysses Sequence Parallelism (SP), a technique from Snowflake AI Research for training LLMs on sequences up to millions of tokens by distributing attention computation across GPUs.

• •Standard transformer attention scales quadratically O(n²) in memory and compute, making 32k+ token training infeasible on a single GPU even with FlashAttention
• •Ulysses shards input sequences across P GPUs, uses all-to-all communication to redistribute by attention heads, so each GPU computes attention for a subset of heads independently
• •Communication cost is O(n·d/P) per GPU — P times lower than Ring Attention's O(n·d), with lower latency due to a single collective step vs. P-1 sequential hops
• •Integrated into Hugging Face Accelerate via ParallelismConfig, Transformers Trainer via TrainingArguments.parallelism_config, and TRL's SFTTrainer with automatic dataloader wrapping and loss aggregation