Ulysses Sequence Parallelism: Training with Million-Token Contexts

Ulysses Sequence Parallelism (SP) is a method developed by Snowflake AI Research as part of the Arctic Long Sequence Training protocol. SP addresses the quadratic memory growth of transformer attention by sharding the input sequence across GPUs and partitioning attention heads, allowing each GPU to compute attention for a subset of heads on a local sequence chunk. The approach requires two all‑to‑all collective operations per attention layer, with a communication volume of 2×(sequence_length/num_gpus)×(head_dimension) per GPU. This design trades sequence locality for head locality, enabling efficient parallelization with low communication overhead. SP is integrated into the Hugging Face ecosystem through Accelerate’s ParallelismConfig, the Transformers Trainer via TrainingArguments.parallelism_config, and TRL’s SFTTrainer, which automatically handles loss aggregation and pre‑shifted labels. The method supports FlashAttention backends and can bypass SP operations during inference. Experiments training Qwen3‑4B on the Gutenberg English dataset with H100 80GB GPUs, DeepSpeed ZeRO‑3, and flash‑attn2 demonstrate scalable long‑context training while maintaining loss equivalence with data‑parallel baselines. The technique requires the number of attention heads to be at least the sequence‑parallel size and that the sequence length be divisible by the SP size. Overall, Ulysses SP enables training on million‑token contexts that would otherwise exceed single‑GPU memory limits, expanding the practical use of large language models for complex, long‑document tasks.