Ulysses Sequence Parallelism for Million-Token Context Training

Hugging Face has integrated Ulysses Sequence Parallelism (part of the Arctic Long Sequence Training protocol from Snowflake AI Research) into the Accelerate, Transformers, and TRL libraries. This integration allows developers to train large language models (LLMs) on sequences of hundreds of thousands or millions of tokens by distributing attention computation across multiple GPUs to overcome the quadratic memory scaling of the attention mechanism.

The Challenge of Long Sequence Training

Standard transformer attention scales quadratically ($O(n^2)$) in both FLOPs and memory relative to sequence length $n$. While FlashAttention reduces memory usage to $O(n)$, the $O(n^2)$ compute requirement remains. For sequences exceeding 32k tokens, training typically exceeds the memory capacity of a single GPU, necessitating a method to split the sequence itself across multiple devices rather than relying solely on data parallelism.

How Ulysses Sequence Parallelism Works

Ulysses Sequence Parallelism (SP) distributes the attention computation by partitioning both the sequence dimension and the attention heads across GPUs. The process follows these steps:

  1. Sequence Sharding: The input sequence is split across $P$ GPUs, where each GPU holds a local chunk of tokens.
  2. QKV Projection: Each GPU computes query, key, and value projections for its local chunk.
  3. All-to-All Communication: An all-to-all collective operation redistributes data so each GPU holds all sequence positions but only for a subset of attention heads.
  4. Local Attention: GPUs compute attention for their assigned heads using FlashAttention or SDPA.
  5. All-to-All Communication: A second all-to-all operation returns the data to the sequence-sharded format.
  6. Output Projection: Each GPU computes the output projection for its local sequence chunk.

Communication Complexity

Ulysses requires two all-to-all operations per attention layer with a communication volume of $O(n · d / P)$ per GPU (where $n$ is sequence length, $d$ is hidden dimension, and $P$ is parallelism degree). This is more efficient than Ring Attention, which communicates $O(n · d)$ per GPU and serializes transfers over $P-1$ hops.

Ecosystem Integration

Accelerate

Accelerate implements Ulysses through the ParallelismConfig class and DeepSpeed integration. Key parameters include sp_size (number of GPUs for sequence parallelism) and sp_backend (which must be set to "deepspeed"). When accelerator.prepare() is called, the system automatically registers the model with UlyssesSPAttentionHF and wraps the dataloader with UlyssesSPDataLoaderAdapter.

Transformers Trainer

The Transformers Trainer handles Ulysses integration via TrainingArguments.parallelism_config. It automates dataloader wrapping, sequence sharding, and weighted loss aggregation, ensuring that gradients are correct even when tokens are unevenly distributed across ranks.

TRL SFTTrainer

TRL's SFTTrainer adds optimizations for supervised fine-tuning, such as the packing feature to reduce padding waste. It requires pad_to_multiple_of to equal the sp_size to ensure sequence divisibility. The SFTTrainer also automatically manages pre-shifted labels when Ulysses is enabled.

Comparing Ulysses and Ring Attention

Aspect Ulysses (DeepSpeed) Ring Attention (FSDP2)
Parallelism Method Attention head partitioning Ring-based KV exchange
Backend DeepSpeed ZeRO PyTorch FSDP2
Attention Support FlashAttention 2/3, SDPA SDPA only
Communication Two all-to-alls per layer P2P ring communication
Comm volume per GPU $O(\text{total_seq} \times \text{hidden} / \text{sp_size})$ $O(\text{total_seq} \times \text{hidden})$
Num Head Constraint num_heads >= sp_size None

Performance Benchmarks

Hugging Face benchmarked Ulysses SP using Qwen3-4B on the Gutenberg English dataset with H100 80GB GPUs.

Memory Reduction

Using SP=4 reduces per-GPU memory by 3.3x at the same sequence length. This enables scaling from a baseline of 8K tokens (DP=4) to 96K tokens (SP=4) while remaining within the 80GB memory limit. At 128K tokens, the configuration reached an Out-of-Memory (OOM) state.

Throughput

Throughput increases as sequence length grows because the quadratic attention computation dominates communication overhead. At 64K tokens, SP=4 achieved 13,396 tokens/second, which is 3.7x the throughput of the 8K baseline.

Best Practices for Implementation

  • Sequence Divisibility: Ensure sequence length is divisible by sp_size using pad_to_multiple_of.
  • Attention Backends: Use FlashAttention 2 for Ampere GPUs and FlashAttention 3 for Hopper GPUs.
  • Memory Optimization: Combine Ulysses with DeepSpeed ZeRO Stage 3 and use the environment variable PYTORCH_ALLOC_CONF=expandable_segments:True to reduce fragmentation.
  • 2D Parallelism: Balance sp_size and dp_shard_size based on GPU count to optimize for either maximum sequence length or higher throughput.
  • Additional Kernels: Use Liger-Kernel's FusedLinearCrossEntropy and TiledMLP to further reduce working memory during loss calculation and large matrix operations.

Sources