Ulysses Sequence Parallelism: Training with Million-Token Contexts

Training large language models on long sequences has change into essential for constructing capable AI systems. As models are increasingly used for tasks like document evaluation, code understanding, complex reasoning, and RAG workloads, the necessity to process sequences of a whole bunch of hundreds—and even thousands and thousands—of tokens has grown dramatically. To place this in perspective, a mean book is roughly 250k tokens, so training on multi-document contexts or book-length inputs requires handling sequences well beyond what matches on a single GPU. Nevertheless, training with such long contexts presents significant memory challenges: the eye computation scales quadratically with sequence length, quickly exceeding GPU memory for contexts beyond tens of hundreds of tokens.

Ulysses Sequence Parallelism (a part of the Arctic Long Sequence Training (ALST) protocol from Snowflake AI Research) provides a chic solution by distributing the eye computation across multiple GPUs through attention head parallelism. On this post, we’ll explore how Ulysses works and the way it has been integrated across the Hugging Face ecosystem—from Speed up to the Transformers Trainer and TRL’s SFTTrainer.

Contents

The Challenge of Long Sequence Training

The eye mechanism in transformers scales quadratically with sequence length. For a sequence of length $nn$, standard attention requires $O(n^2)O(n^2)$

Consider these scenarios where long-context training is crucial:

• Document understanding: Processing entire books, legal documents, or research papers
• Code evaluation: Understanding large codebases with multiple interconnected files
• Reasoning tasks: Models that “think” step-by-step may generate hundreds of tokens during inference
• Retrieval-augmented generation: Incorporating many retrieved passages into the context

Traditional data parallelism doesn’t help here—each GPU still must process the total sequence contained in the attention block. We want a solution to split the sequence itself across multiple devices.

How Ulysses Works

Ulysses Sequence Parallelism (SP), introduced within the DeepSpeed Ulysses paper, takes a clever approach: along with splitting on the sequence dimension, it also partitions the eye heads across GPUs.

Here’s how it really works:

1. Sequence Sharding: The input sequence is split along the sequence dimension across $PP$ GPUs. Each GPU $ii$ holds tokens $[i\cdot n\mathrm{/}P,(i+1)\cdot n\mathrm{/}P)[i\; cdot\; n/P,\; (i+1)\; cdot\; n/P)$

2. QKV Projection: Each GPU computes the query, key, and value projections for its local sequence chunk.

3. All-to-All Communication: An all-to-all collective operation redistributes the data so that each GPU holds all sequence positions after the projections, but only for a subset of attention heads.

4. Local Attention: Each GPU computes attention for its assigned heads using standard attention mechanisms (FlashAttention or SDPA).

5. All-to-All Communication: Another all-to-all operation reverses the redistribution, returning to sequence-sharded format.

6. 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.

Communication Complexity

Ulysses requires two all-to-all operations per attention layer, with total communication volume of $O(n\cdot d\mathrm{/}P)O(n\; cdot\; d\; /\; P)$

• $nn$ is the sequence length
• $dd$ is the hidden dimension
• $PP$ is the parallelism degree

Ring Attention communicates $O(n\cdot d)O(n\; cdot\; d)$

Integration with Accelerate

Accelerate provides the foundation for Ulysses sequence parallelism through its ParallelismConfig class and DeepSpeed integration.

Configuration

from accelerate import Accelerator
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,  
    dp_shard_size=1,  
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=None,  
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",  
    ),
)

accelerator = Accelerator(parallelism_config=parallelism_config)

Key Parameters

Using the Accelerator

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)


model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

The prepare() call:

1. Registers the model with DeepSpeed’s UlyssesSPAttentionHF
2. Wraps the dataloader with UlyssesSPDataLoaderAdapter to handle sequence sharding
3. Automatically injects shift_labels for correct loss computation

Loss Aggregation

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()

    
    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)

    
    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).

Each Ulysses and Ring Attention use position_ids as an alternative of attention_mask for causal masking during training. A 4D attention mask at these sequence lengths could be just as prohibitive as the eye scores themselves—at 128k tokens, that is one other ~1TB tensor. Position IDs achieve the identical causal behavior with $O(n)O(n)$ memory as an alternative of $O(n^2)O(n^2)$

Integration with Transformers Trainer

The Transformers Trainer provides seamless Ulysses integration through TrainingArguments.parallelism_config. It handles all of the SP-specific details mechanically—dataloader wrapping, sequence sharding, and loss aggregation—so you do not need to write down any of the custom loss code shown above.

Configuration

Just pass the identical parallelism_config from above into TrainingArguments:

from transformers import TrainingArguments

training_args = TrainingArguments(
    parallelism_config=parallelism_config,  
    per_device_train_batch_size=1,
)

What the Trainer Handles Routinely

1. Dataloader Wrapping: After model preparation, the Trainer wraps the dataloader with UlyssesSPDataLoaderAdapter

2. Loss Computation: The compute_loss method detects SP mode and routes to specialized _deepspeed_sp_compute_loss which handles:

   • Gathering losses across SP ranks
   • Computing valid token counts per rank
   • Weighted loss aggregation

3. Batch Size Calculation: The effective data parallel world size accounts for SP:

   dp_world_size = world_size // sp_size
   

4. Dataloader Length Adjustment: Training step calculations are adjusted for SP’s effect on iteration count

Launch Command

Use an speed up config file or command-line arguments:

speed up launch 
    --config_file deepspeed_ulysses.yaml 
    train.py 
    --per_device_train_batch_size 1

Integration with TRL SFTTrainer

TRL’s SFTTrainer builds on the Transformers Trainer and adds specific optimizations for supervised fine-tuning with long sequences.

Configuration

from trl import SFTConfig, SFTTrainer
from speed up.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=2,
    dp_shard_size=2,  
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

training_args = SFTConfig(
    ...,
    parallelism_config=parallelism_config,
    max_length=32768,
    pad_to_multiple_of=2,  
    per_device_train_batch_size=1,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

Key SFTConfig Parameters for Ulysses

Speed up Config File

Create alst_ulysses_4gpu.yaml:

compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
mixed_precision: bf16
num_processes: 4
deepspeed_config:
  zero_stage: 3
  seq_parallel_communication_data_type: bf16
parallelism_config:
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_dp_shard_size: 2
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: flash_attention_2

Complete Training Command

speed up launch --config_file alst_ulysses_4gpu.yaml 
    trl/scripts/sft.py 
    --model_name_or_path meta-llama/Llama-3.1-8B 
    --dataset_name trl-lib/Capybara 
    --max_length 32768 
    --packing 
    --pad_to_multiple_of 2 
    --per_device_train_batch_size 1

Shift Labels Handling

The SFTTrainer mechanically handles pre-shifted labels when Ulysses is enabled:

labels = inputs["labels"] if "shift_labels" not in inputs else None


if "shift_labels" in inputs:
    shift_logits = outputs.logits.contiguous()
    shift_labels = inputs["shift_labels"]
else:
    shift_logits = outputs.logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

Comparing Ulysses and Ring Attention

Each Ulysses and Ring Attention enable long-context training, but they’ve different characteristics:

When to Select Ulysses vs Ring Attention

Since switching between the 2 only requires changing the speed up config, we recommend trying each and comparing performance and memory usage in your specific setup. The principal constraint is that Ulysses requires num_heads >= sp_size, while Ring Attention has no such limitation.

Best Practices

1. Sequence Length Divisibility

At all times ensure your sequence length is divisible by sp_size:

training_args = SFTConfig(
    pad_to_multiple_of=4,  
    max_length=32768,  
)

2. Use Flash Attention

Flash Attention 2 provides cleaner output and higher performance than SDPA:

parallelism_config = ParallelismConfig(
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_attn_implementation="flash_attention_2",
    ),
)

Use Flash Attention 3 for Hopper and look out for Flash Attention 4 release for Blackwell (FA2 on Blackwell is sort of slow).

3. Mix with DeepSpeed ZeRO

For very large models, mix Ulysses with ZeRO Stage 3:

deepspeed_config:
  zero_stage: 3
  offload_optimizer:
    device: cpu

If the model is large, you possibly can offload the params as well by adding to the above:

  offload_param:
    device: cpu

5. Use memory fragmentation-friendly PyTorch allocator

This environment variable will allow for an extended sequence length:

export PYTORCH_ALLOC_CONF=expandable_segments:True

6. 2D Parallelism Configuration

Balance SP and DP to your GPU count:

Remember: dp_replicate_size × dp_shard_size × sp_size = num_processes

7. Liger-Kernel

If your required model architecture is supported by Liger-Kernel, it’s fully compatible with Ulysses SP and could be enabled with a single flag:

training_args = SFTConfig(
    use_liger_kernel=True,
)

The principal memory saving comes from FusedLinearCrossEntropy which avoids materializing the total logits tensor during loss calculation. The savings grow with longer sequences where the logits tensor is larger.

Moreover, you possibly can enable TiledMLP to further extend sequence length — like FusedLinearCrossEntropy, it saves working memory by tiling large matrix operations.

8. Token Distribution Across Ranks

You need not worry about manually balancing tokens across SP ranks—the loss aggregation code handles uneven distributions gracefully (including ranks with zero valid tokens). With random batching over a fairly sized dataset, the distribution evens out statistically over training.

Benchmarks

To quantify the advantages of Ulysses SP, we trained Qwen3-4B on the Gutenberg English streaming dataset using TRL’s SFTTrainer. All experiments ran on H100 80GB GPUs with DeepSpeed ZeRO-3, CPU optimizer offloading, gradient checkpointing, and flash-attn2 as the eye backend.

Setup

The benchmark runs within the table above use the identical global batch size (8 micro-batches), cosine learning-rate schedule, and seed, so those benchmark loss curves are directly comparable.

Loss Curve Matching Diagnostics (4 GPU)

To confirm SP-vs-DP loss equivalence, we ran controlled 4-GPU A/B experiments with an identical seed, model, optimizer, learning-rate schedule, and data order.

Methodology for Fair DP vs SP Comparison

Compared setups:

• DP=4, SP=1, GAS=1 (baseline)
• DP=1, SP=4, GAS=4 (Ulysses SP)

For fair comparison, GAS must scale with SP:

• Ulysses SP splits the sequence across SP ranks, so each SP rank sees roughly 1/SP of the sequence tokens per micro-step.
• If GAS is unchanged, each optimizer step in SP aggregates fewer total tokens than the DP baseline.
• Setting GAS=SP keeps effective tokens per optimizer step matched:
  • DP tokens/step: dp_world_size * micro_batch * seq_len * GAS = 4 * B * L * 1
  • SP tokens/step: dp_world_size * micro_batch * (L/SP) * GAS * SP_ranks = 1 * B * (L/4) * 4 * 4 = 4 * B * L

Measured over 20 steps on 4 GPUs in controlled equivalence harnesses:

Takeaway: under matched token budget, SP and non-SP match on canonical token-normalized loss. The remaining difference is in trainer-reported logging (loss), not within the underlying cross-entropy objective.

Memory Reduction

At 8K tokens, DP=4 and SP=4 use nearly the identical memory per GPU (~22 GB with ZeRO-3). The advantage of SP is that it enables scaling to for much longer sequences: at 96K tokens (12x longer), peak memory is 66 GB — still throughout the H100’s 80 GB capability. At 128K, the model OOMs, establishing the sensible limit for this configuration. DP=4 without SP cannot scale beyond 8K for this model.

Throughput

At the identical sequence length (8K), SP=4 has comparable throughput to the single-GPU baseline — the all-to-all communication overhead is minimal on NVLink-connected GPUs. The true profit comes from longer sequences: as sequence length grows, the quadratic attention computation dominates over communication and other overheads, making each training step increasingly compute-efficient. Each step also processes proportionally more tokens, so throughput scales with sequence length. At 64K, SP=4 processes 13,396 tokens/second — 3.7x the baseline.

These results use only 4 GPUs with SP=4. With 8 GPUs (SP=8), you possibly can push to even longer sequences — as much as 256K+ tokens — or use 2D parallelism (SP=4, DP=2) to mix long-context training with data-parallel throughput.

Requirements

• HF Speed up: deepspeed>=0.18.1 speed up>=1.12
• HF Trainer: deepspeed>=0.18.1 speed up>=1.12 transformers>=5.0
• HF TRL: deepspeed>=0.18.1 speed up>=1.12 transformers>=5.0 trl>=0.18.0

Use flash_attention_2 for Ampere GPUs, or flash_attention_3 for Hopper GPUs. Wait for flash_attention_4 on Blackwell 🕰.

Resources

Documentation

Examples

Papers

Related Blog Posts