Ulysses Sequence Parallelism: Training with Million-Token Contexts

Training large language models on long sequences has become essential for building capable AI systems. As models are increasingly used for tasks like document analysis, code understanding, complex reasoning, and RAG workloads, the need to process sequences of hundreds of thousands—or even millions—of tokens has grown dramatically. To put this in perspective, an average book is roughly 250k tokens, so training on multi-document contexts or book-length inputs requires handling sequences well beyond what fits on a single GPU. However, training with such long contexts presents significant memory challenges: the attention computation scales quadratically with sequence length, quickly exceeding GPU memory for contexts beyond tens of thousands of tokens.

Ulysses Sequence Parallelism (part of the Arctic Long Sequence Training (ALST) protocol from Snowflake AI Research) provides an elegant solution by distributing the attention computation across multiple GPUs through attention head parallelism. In this post, we'll explore how Ulysses works and how it's been integrated across the Hugging Face ecosystem—from Accelerate to the Transformers Trainer and TRL's SFTTrainer.

• The Challenge of Long Sequence Training
• How Ulysses Works
• Integration with Accelerate
• Integration with Transformers Trainer
• Integration with TRL's SFTTrainer
• Comparing Ulysses and Ring Attention
• Best Practices
• Benchmarks
• Resources

The attention mechanism in transformers scales quadratically with sequence length. For a sequence of length , standard attention requires FLOPs and memory to compute and store the attention score matrix. Optimized implementations like FlashAttention reduce the memory to by tiling the computation and never materializing the full attention matrix—but the compute remains. For very long sequences (32k+ tokens), even with FlashAttention, training still pushes the limits of single-GPU memory.

Consider these scenarios where long-context training is essential:

Document understanding: Processing entire books, legal documents, or research papersCode analysis: Understanding large codebases with multiple interconnected filesReasoning tasks: Models that "think" step-by-step may generate thousands of tokens during inferenceRetrieval-augmented generation: Incorporating many retrieved passages into the context

Traditional data parallelism doesn't help here—each GPU still needs to process the full sequence inside the attention block. We need a way to split the sequence itself across multiple devices.

Ulysses Sequence Parallelism (SP), introduced in the DeepSpeed Ulysses paper, takes a clever approach: in addition to splitting on the sequence dimension, it also partitions the attention heads across GPUs.

Here's how it works:

Sequence Sharding: The input sequence is split along the sequence dimension across GPUs. Each GPU holds tokens .QKV Projection: Each GPU computes the query, key, and value projections for its local sequence chunk.All-to-All Communication: An all-to-all collective operation redistributes the data so that each GPU holdsallsequence positions after the projections, but only for a subset of attention heads.Local Attention: Each GPU computes attention for its assigned heads using standard attention mechanisms (FlashAttention or SDPA).All-to-All Communication: Another all-to-all operation reverses the redistribution, returning to sequence-sharded format.Output Projection: Each GPU computes the output projection for its local sequence chunk.

The key insight is that attention heads are independent—each head can be computed separately. By trading sequence locality for head locality, Ulysses enables efficient parallelization with relatively low communication overhead.

Ulysses requires two all-to-all operations per attention layer, with total communication volume of per GPU, where:

• is the sequence length
• is the hidden dimension
• is the parallelism degree

Ring Attention communicates per GPU — a factor of more — via sequential point-to-point transfers around the ring. Ulysses also benefits from lower latency because all-to-all can exploit full bisectional bandwidth in a single collective step, whereas Ring Attention serializes over hops.

Accelerate provides the foundation for Ulysses sequence parallelism through its ParallelismConfig

class and DeepSpeed integration.

from accelerate import Accelerator
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig
parallelism_config = ParallelismConfig(
sp_backend="deepspeed",
sp_size=4, # Split across 4 GPUs
dp_shard_size=1, # Must satisfy: dp_replicate × dp_shard × sp_size = num_processes
sp_handler=DeepSpeedSequenceParallelConfig(
sp_seq_length=None, # None for variable-length sequences
sp_seq_length_is_variable=True,
sp_attn_implementation="flash_attention_2", # or "sdpa"
),
)
accelerator = Accelerator(parallelism_config=parallelism_config)

| Parameter | Description | |---|---| sp_size | Number of GPUs for sequence parallelism | sp_backend | Must be "deepspeed" for Ulysses | sp_seq_length_is_variable | Set to True for varying sequence lengths across batches | sp_attn_implementation | "flash_attention_2" , "flash_attention_3" , or "sdpa" |

When you call accelerator.prepare()

, Ulysses is automatically set up:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
# This registers the model with Ulysses and wraps the dataloader
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

The prepare()

call:

• Registers the model with DeepSpeed's UlyssesSPAttentionHF

• Wraps the dataloader with UlyssesSPDataLoaderAdapter

to handle sequence sharding - Automatically injects shift_labels

for correct loss computation

With Ulysses, each GPU computes loss on different parts of the sequence. The losses must be aggregated properly, weighted by the number of valid tokens per rank. If you're using the Transformers Trainer

or TRL's SFTTrainer

, this is handled automatically—the code below is only needed when writing a custom Accelerate training loop:

sp_size = parallelism_config.sp_size
if sp_size > 1:
from deepspeed.utils import groups
sp_group = groups._get_sequence_parallel_group()
# Gather losses and token counts from all SP ranks
losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
good_tokens = (batch["shift_labels"] != -100).view(-1).sum()
good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)
# Weighted aggregation
total_loss = sum(
losses_per_rank[i] * good_tokens_per_rank[i]
for i in range(sp_size)
if good_tokens_per_rank[i] > 0
)
loss = total_loss / max(sum(good_tokens_per_rank), 1)
accelerator.backward(loss)

The weighted loss aggregation ensures correct gradients when tokens are unevenly distributed across ranks (e.g., when some ranks contain only padding or masked out prompt tokens).

Both Ulysses and Ring Attention use

position_ids

instead ofattention_mask

for causal masking during training. A 4D attention mask at these sequence lengths would be just as prohibitive as the attention scores themselves—at 128k tokens, that's another ~1TB tensor. Position IDs achieve the same causal behavior with memory instead of . During evaluation/inference, DeepSpeed's SP attention layer can bypass the SP operations entirely (viadisable_in_eval

) and fall back to the model's default attention implementation.

The Transformers Trainer

provides seamless Ulysses integration through TrainingArguments.parallelism_config

. It handles all the SP-specific details automatically—dataloader wrapping, sequence sharding, and loss aggregation—so you don't need to write any of the custom loss code shown above.

Just pass the same parallelism_config

from ab