{
  "@context": "https://schema.org",
  "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-03-09-ulysses-sequence-parallelism-training-with-million-token/#webpage",
  "@type": [
    "WebPage",
    "ArchiveComponent"
  ],
  "articleBody": "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.\n\nUlysses Sequence Parallelism (part of the [Arctic Long Sequence Training (ALST) protocol](https://huggingface.co/papers/2506.13996) 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.\n\n## Contents\n\n- [The Challenge of Long Sequence Training](https://huggingface.co/blog/ulysses-sp#the-challenge-of-long-sequence-training)\n- [How Ulysses Works](https://huggingface.co/blog/ulysses-sp#how-ulysses-works)\n- [Integration with Accelerate](https://huggingface.co/blog/ulysses-sp#integration-with-accelerate)\n- [Integration with Transformers Trainer](https://huggingface.co/blog/ulysses-sp#integration-with-transformers-trainer)\n- [Integration with TRL's SFTTrainer](https://huggingface.co/blog/ulysses-sp#integration-with-trl-sfttrainer)\n- [Comparing Ulysses and Ring Attention](https://huggingface.co/blog/ulysses-sp#comparing-ulysses-and-ring-attention)\n- [Best Practices](https://huggingface.co/blog/ulysses-sp#best-practices)\n- [Benchmarks](https://huggingface.co/blog/ulysses-sp#benchmarks)\n- [Resources](https://huggingface.co/blog/ulysses-sp#resources)\n\n## The Challenge of Long Sequence Training\n\nThe attention mechanism in transformers scales quadratically with sequence length. For a sequence of length n n, standard attention requires O(n2) O(n^2) FLOPs and O(n2) O(n^2) memory to compute and store the attention score matrix. Optimized implementations like [FlashAttention](https://huggingface.co/papers/2205.14135) reduce the **memory** to O(n) O(n) by tiling the computation and never materializing the full attention matrix—but the O(n2) O(n^2) compute remains. For very long sequences (32k+ tokens), even with FlashAttention, training still pushes the limits of single-GPU memory.\n\nConsider these scenarios where long-context training is essential:\n\n- **Document understanding**: Processing entire books, legal documents, or research papers\n- **Code analysis**: Understanding large codebases with multiple interconnected files\n- **Reasoning tasks**: Models that \"think\" step-by-step may generate thousands of tokens during inference\n- **Retrieval-augmented generation**: Incorporating many retrieved passages into the context\n\nTraditional 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.\n\n## How Ulysses Works\n\nUlysses Sequence Parallelism (SP), introduced in the [DeepSpeed Ulysses paper](https://huggingface.co/papers/2309.14509), takes a clever approach: in addition to splitting on the sequence dimension, it also partitions the attention heads across GPUs.\n\n![Ulysses Sequence Parallelism Overview](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/ulysses_overview.png)\\\nUlysses splits input sequences along the sequence dimension and uses all-to-all communication to exchange key-value pairs, enabling each GPU to compute a subset of attention heads. (***Source: [Snowflake Engineering Blog](https://www.snowflake.com/en/engineering-blog/ulysses-low-latency-llm-inference/)***)\n\nHere's how it works:\n\n1. **Sequence Sharding**: The input sequence is split along the sequence dimension across P P GPUs. Each GPU i i holds tokens \\[i⋅n/P,(i+1)⋅n/P) \\[i \\\\cdot n/P, (i+1) \\\\cdot n/P).\n\n2. **QKV Projection**: Each GPU computes the query, key, and value projections for its local sequence chunk.\n\n3. **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.\n\n4. **Local Attention**: Each GPU computes attention for its assigned heads using standard attention mechanisms (FlashAttention or SDPA).\n\n5. **All-to-All Communication**: Another all-to-all operation reverses the redistribution, returning to sequence-sharded format.\n\n6. **Output Projection**: Each GPU computes the output projection for its local sequence chunk.\n\nThe 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.\n\n### Communication Complexity\n\nUlysses requires two all-to-all operations per attention layer, with total communication volume of O(n⋅d/P) O(n \\\\cdot d / P) per GPU, where:\n\n- n n is the sequence length\n- d d is the hidden dimension\n- P P is the parallelism degree\n\nRing Attention communicates O(n⋅d) O(n \\\\cdot d) per GPU — a factor of P P more — via P−1 P-1 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 P−1 P-1 hops.\n\n## Integration with Accelerate\n\nAccelerate provides the foundation for Ulysses sequence parallelism through its [`ParallelismConfig`](https://huggingface.co/docs/accelerate/concept_guides/sequence_parallelism#alstulysses-sp-backend-configuration) class and DeepSpeed integration.\n\n### Configuration\n\n```python\nfrom accelerate import Accelerator\nfrom accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig\n\nparallelism_config = ParallelismConfig(\n    sp_backend=\"deepspeed\",\n    sp_size=4,  # Split across 4 GPUs\n    dp_shard_size=1,  # Must satisfy: dp_replicate × dp_shard × sp_size = num_processes\n    sp_handler=DeepSpeedSequenceParallelConfig(\n        sp_seq_length=None,  # None for variable-length sequences\n        sp_seq_length_is_variable=True,\n        sp_attn_implementation=\"flash_attention_2\",  # or \"sdpa\"\n    ),\n)\n\naccelerator = Accelerator(parallelism_config=parallelism_config)\n```\n\n### Key Parameters\n\n| Parameter                   | Description                                               |\n| --------------------------- | --------------------------------------------------------- |\n| `sp_size`                   | Number of GPUs for sequence parallelism                   |\n| `sp_backend`                | Must be `\"deepspeed\"` for Ulysses                         |\n| `sp_seq_length_is_variable` | Set to `True` for varying sequence lengths across batches |\n| `sp_attn_implementation`    | `\"flash_attention_2\"`, `\"flash_attention_3\"`, or `\"sdpa\"` |\n\n### Using the Accelerator\n\nWhen you call `accelerator.prepare()`, Ulysses is automatically set up:\n\n```python\nfrom transformers import AutoModelForCausalLM\n\nmodel = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-3.1-8B\")\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)\n\n# This registers the model with Ulysses and wraps the dataloader\nmodel, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)\n```\n\nThe `prepare()` call:\n\n1. Registers the model with DeepSpeed's `UlyssesSPAttentionHF`\n2. Wraps the dataloader with `UlyssesSPDataLoaderAdapter` to handle sequence sharding\n3. Automatically injects `shift_labels` for correct loss computation\n\n### Loss Aggregation\n\nWith 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:\n\n```python\nsp_size = parallelism_config.sp_size\nif sp_size > 1:\n    from deepspeed.utils import groups\n\n    sp_group = groups._get_sequence_parallel_group()\n\n    # Gather losses and token counts from all SP ranks\n    losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)\n    good_tokens = (batch[\"shift_labels\"] != -100).view(-1).sum()\n    good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)\n\n    # Weighted aggregation\n    total_loss = sum(\n        losses_per_rank[i] * good_tokens_per_rank[i]\n        for i in range(sp_size)\n        if good_tokens_per_rank[i] > 0\n    )\n    loss = total_loss / max(sum(good_tokens_per_rank), 1)\n\naccelerator.backward(loss)\n```\n\n> 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).\n\n> Both Ulysses and Ring Attention use `position_ids` instead of `attention_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 O(n) O(n) memory instead of O(n2) O(n^2). During evaluation/inference, DeepSpeed's SP attention layer can bypass the SP operations entirely (via `disable_in_eval`) and fall back to the model's default attention implementation.\n\n## Integration with Transformers Trainer\n\nThe Transformers [`Trainer`](https://huggingface.co/docs/transformers/main_classes/trainer) provides seamless Ulysses integration through [`TrainingArguments.parallelism_config`](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.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.\n\n### Configuration\n\nJust pass the same `parallelism_config` from above into `TrainingArguments`:\n\n```python\nfrom transformers import TrainingArguments\n\ntraining_args = TrainingArguments(\n    parallelism_config=parallelism_config,  # same ParallelismConfig as above\n    per_device_train_batch_size=1,\n)\n```\n\n### What the Trainer Handles Automatically\n\n1. **Dataloader Wrapping**: After model preparation, the Trainer wraps the dataloader with `UlyssesSPDataLoaderAdapter`\n\n2. **Loss Computation**: The `compute_loss` method detects SP mode and routes to specialized `_deepspeed_sp_compute_loss` which handles:\n\n   - Gathering losses across SP ranks\n   - Computing valid token counts per rank\n   - Weighted loss aggregation\n3. **Batch Size Calculation**: The effective data parallel world size accounts for SP:\n\n   ```python\n   dp_world_size = world_size // sp_size\n   ```\n\n4. **Dataloader Length Adjustment**: Training step calculations are adjusted for SP's effect on iteration count\n\n### Launch Command\n\nUse an accelerate config file or command-line arguments:\n\n```bash\naccelerate launch \\\n    --config_file deepspeed_ulysses.yaml \\\n    train.py \\\n    --per_device_train_batch_size 1\n```\n\n## Integration with TRL SFTTrainer\n\nTRL's [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) builds on the Transformers Trainer and adds specific optimizations for supervised fine-tuning with long sequences.\n\n### Configuration\n\n```python\nfrom trl import SFTConfig, SFTTrainer\nfrom accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig\n\nparallelism_config = ParallelismConfig(\n    sp_backend=\"deepspeed\",\n    sp_size=2,\n    dp_shard_size=2,  # 2D parallelism: SP × DP = 4 GPUs\n    sp_handler=DeepSpeedSequenceParallelConfig(\n        sp_seq_length_is_variable=True,\n        sp_attn_implementation=\"flash_attention_2\",\n    ),\n)\n\ntraining_args = SFTConfig(\n    ...,\n    parallelism_config=parallelism_config,\n    max_length=32768,\n    pad_to_multiple_of=2,  # Must equal sp_size\n    per_device_train_batch_size=1,\n)\n\ntrainer = SFTTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=dataset,\n)\ntrainer.train()\n```\n\n### Key SFTConfig Parameters for Ulysses\n\n| Parameter            | Description                                                                                  |\n| -------------------- | -------------------------------------------------------------------------------------------- |\n| `pad_to_multiple_of` | Must equal `sp_size` to ensure sequence divisibility                                         |\n| `max_length`         | Global sequence length (before splitting across GPUs)                                        |\n| `packing`            | Works well with SP — packing reduces padding waste, especially for variable-length sequences |\n\n### Accelerate Config File\n\nCreate `alst_ulysses_4gpu.yaml`:\n\n```yaml\ncompute_environment: LOCAL_MACHINE\ndistributed_type: DEEPSPEED\nmixed_precision: bf16\nnum_processes: 4\ndeepspeed_config:\n  zero_stage: 3\n  seq_parallel_communication_data_type: bf16\nparallelism_config:\n  parallelism_config_sp_size: 2\n  parallelism_config_sp_backend: deepspeed\n  parallelism_config_dp_shard_size: 2\n  parallelism_config_sp_seq_length_is_variable: true\n  parallelism_config_sp_attn_implementation: flash_attention_2\n```\n\n### Complete Training Command\n\n```bash\naccelerate launch --config_file alst_ulysses_4gpu.yaml \\\n    trl/scripts/sft.py \\\n    --model_name_or_path meta-llama/Llama-3.1-8B \\\n    --dataset_name trl-lib/Capybara \\\n    --max_length 32768 \\\n    --packing \\\n    --pad_to_multiple_of 2 \\\n    --per_device_train_batch_size 1\n```\n\n### Shift Labels Handling\n\nThe SFTTrainer automatically handles pre-shifted labels when Ulysses is enabled:\n\n```python\n# When using SP, labels are pre-shifted by the dataloader adapter\n# The trainer detects this and uses shift_labels directly\nlabels = inputs[\"labels\"] if \"shift_labels\" not in inputs else None\n\n# Loss computation uses the pre-shifted labels\nif \"shift_labels\" in inputs:\n    shift_logits = outputs.logits.contiguous()\n    shift_labels = inputs[\"shift_labels\"]\nelse:\n    shift_logits = outputs.logits[..., :-1, :].contiguous()\n    shift_labels = labels[..., 1:].contiguous()\n```\n\n## Comparing Ulysses and Ring Attention\n\nBoth Ulysses and Ring Attention enable long-context training, but they have different characteristics:\n\n| Aspect                    | Ulysses (DeepSpeed)               | Ring Attention (FSDP2) |\n| ------------------------- | --------------------------------- | ---------------------- |\n| **Parallelism Method**    | Attention head partitioning       | Ring-based KV exchange |\n| **Backend**               | DeepSpeed ZeRO                    | PyTorch FSDP2          |\n| **Attention Support**     | FlashAttention 2/3, SDPA          | SDPA only              |\n| **Communication**         | Two `all-to-all`s per layer       | P2P ring communication |\n| **Comm volume per GPU**   | O(total\\_seq x hidden / sp\\_size) | O(total\\_seq x hidden) |\n| **Sequence Divisibility** | `sp_size`                         | `cp_size * 2`          |\n| **Num Head Constraint**   | `num_heads >= sp_size`            | None                   |\n\n### When to Choose Ulysses vs Ring Attention\n\nSince switching between the two only requires changing the accelerate config, we recommend trying both and comparing performance and memory usage on your specific setup. The main constraint is that Ulysses requires `num_heads >= sp_size`, while Ring Attention has no such limitation.\n\n## Best Practices\n\n### 1. Sequence Length Divisibility\n\nAlways ensure your sequence length is divisible by `sp_size`:\n\n```python\ntraining_args = SFTConfig(\n    pad_to_multiple_of=4,  # For sp_size=4\n    max_length=32768,  # Must be divisible by 4\n)\n```\n\n### 2. Use Flash Attention\n\nFlash Attention 2 provides cleaner output and better performance than SDPA:\n\n```python\nparallelism_config = ParallelismConfig(\n    sp_handler=DeepSpeedSequenceParallelConfig(\n        sp_attn_implementation=\"flash_attention_2\",\n    ),\n)\n```\n\nUse Flash Attention 3 for Hopper and look out for Flash Attention 4 release for Blackwell (FA2 on Blackwell is quite slow).\n\n### 3. Combine with DeepSpeed ZeRO\n\nFor very large models, combine Ulysses with ZeRO Stage 3:\n\n```yaml\ndeepspeed_config:\n  zero_stage: 3\n  offload_optimizer:\n    device: cpu\n```\n\nIf the model is huge, you can offload the params as well by adding to the above:\n\n```yaml\n  offload_param:\n    device: cpu\n```\n\n### 5. Use memory fragmentation-friendly PyTorch allocator\n\nThis environment variable will allow for a longer sequence length:\n\n```bash\nexport PYTORCH_ALLOC_CONF=expandable_segments:True\n```\n\n### 6. 2D Parallelism Configuration\n\nBalance SP and DP for your GPU count:\n\n| GPUs | `sp_size` | `dp_shard_size` | Use Case                                        |\n| ---- | --------- | --------------- | ----------------------------------------------- |\n| 4    | 2         | 2               | Balanced throughput and sequence length         |\n| 4    | 4         | 1               | Maximum sequence length                         |\n| 8    | 2         | 4               | Higher throughput with moderate sequence length |\n| 8    | 4         | 2               | Longer sequences with moderate throughput       |\n\nRemember: `dp_replicate_size × dp_shard_size × sp_size = num_processes`\n\n### 7. Liger-Kernel\n\nIf your desired model architecture is supported by [Liger-Kernel](https://github.com/linkedin/Liger-Kernel), it is fully compatible with Ulysses SP and can be enabled with a single flag:\n\n```python\ntraining_args = SFTConfig(\n    use_liger_kernel=True,\n)\n```\n\nThe main memory saving comes from `FusedLinearCrossEntropy` which avoids materializing the full logits tensor during loss calculation. The savings grow with longer sequences where the logits tensor is larger.\n\nAdditionally, you can enable [`TiledMLP`](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/#tiled-mlp-computation) to further extend sequence length — like `FusedLinearCrossEntropy`, it saves working memory by tiling large matrix operations.\n\n### 8. Token Distribution Across Ranks\n\nYou don't need to 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 reasonably sized dataset, the distribution evens out statistically over training.\n\n## Benchmarks\n\nTo quantify the benefits of Ulysses SP, we trained [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) on the [Gutenberg English](https://huggingface.co/datasets/sedthh/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 attention backend.\n\n### Setup\n\n| Config   | GPUs | SP | DP | Seq Length | Grad Acc | Global Batch |\n| -------- | ---- | -- | -- | ---------- | -------- | ------------ |\n| Baseline | 1    | 1  | 1  | 8K         | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 8K         | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 32K        | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 64K        | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 96K        | 8        | 8            |\n\nThe benchmark runs in the table above use the same global batch size (8 micro-batches), cosine learning-rate schedule, and seed, so those benchmark loss curves are directly comparable.\n\n### Loss Curve Matching Diagnostics (4 GPU)\n\nTo verify SP-vs-DP loss equivalence, we ran controlled 4-GPU A/B experiments with identical seed, model, optimizer, learning-rate schedule, and data order.\n\n#### Methodology for Fair DP vs SP Comparison\n\nCompared setups:\n\n- `DP=4, SP=1, GAS=1` (baseline)\n- `DP=1, SP=4, GAS=4` (Ulysses SP)\n\nFor fair comparison, `GAS` must scale with `SP`:\n\n- Ulysses SP splits the sequence across `SP` ranks, so each SP rank sees roughly `1/SP` of the sequence tokens per micro-step.\n- If `GAS` is unchanged, each optimizer step in SP aggregates fewer total tokens than the DP baseline.\n- Setting `GAS=SP` keeps effective tokens per optimizer step matched:\n  - DP tokens/step: `dp_world_size * micro_batch * seq_len * GAS = 4 * B * L * 1`\n  - SP tokens/step: `dp_world_size * micro_batch * (L/SP) * GAS * SP_ranks = 1 * B * (L/4) * 4 * 4 = 4 * B * L`\n\n![Canonical loss on Gutenberg for DP=4 vs SP=4](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/gutenberg_canonical_loss_dp4_vs_sp4.png)\\\nOn Gutenberg text (20 steps), canonical loss matches within logging precision between `DP=4,SP=1,GAS=1` and `DP=1,SP=4,GAS=4`.\n\nMeasured over 20 steps on 4 GPUs in controlled equivalence harnesses:\n\n| Harness      | Metric        | DP vs SP setting         | Mean abs diff | Max abs diff |\n| ------------ | ------------- | ------------------------ | ------------- | ------------ |\n| `Trainer`    | `loss`        | DP=4, SP=1 vs DP=1, SP=4 | 0.0054        | 0.0131       |\n| `SFTTrainer` | logged `loss` | DP=4, SP=1 vs DP=1, SP=4 | 0.0811        | 0.0812       |\n| `SFTTrainer` | canonical NLL | DP=4, SP=1 vs DP=1, SP=4 | 0.000004      | 0.000005     |\n\nTakeaway: under matched token budget, SP and non-SP match on canonical token-normalized loss. The remaining difference is in trainer-reported logging (`loss`), not in the underlying cross-entropy objective.\n\n### Memory Reduction\n\n![Peak GPU Memory per Rank](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/memory_usage.png)\\\nSP=4 reduces per-GPU memory by 3.3x at the same sequence length, enabling training at up to 96K tokens on 4× H100 80GB. At 128K, the model OOMs.\n\n| Config       | Seq Length | Peak Memory | Notes                             |\n| ------------ | ---------- | ----------- | --------------------------------- |\n| DP=4 (4 GPU) | 8K         | 22.4 GB     | Baseline — no SP                  |\n| SP=4 (4 GPU) | 8K         | 22.8 GB     | Similar memory at same seq length |\n| SP=4 (4 GPU) | 32K        | 35.0 GB     | 4x longer than DP baseline        |\n| SP=4 (4 GPU) | 64K        | 50.5 GB     | 8x longer than DP baseline        |\n| SP=4 (4 GPU) | 96K        | 66.0 GB     | **12x longer than DP baseline**   |\n| SP=4 (4 GPU) | 128K       | OOM         | Exceeds 80 GB limit               |\n\nAt 8K tokens, DP=4 and SP=4 use nearly the same memory per GPU (~22 GB with ZeRO-3). The advantage of SP is that it enables scaling to much longer sequences: at 96K tokens (12x longer), peak memory is 66 GB — still within the H100's 80 GB capacity. At 128K, the model OOMs, establishing the practical limit for this configuration. DP=4 without SP cannot scale beyond 8K for this model.\n\n### Throughput\n\n![Training Throughput](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/throughput.png)\\\nLonger sequences with SP process dramatically more tokens per second. SP=4 at 64K achieves 3.7x the throughput of the baseline.\n\n| Config           | Seq Length | Tokens/s | vs Baseline |\n| ---------------- | ---------- | -------- | ----------- |\n| Baseline (1 GPU) | 8K         | 3,633    | —           |\n| SP=4 (4 GPU)     | 8K         | 3,933    | \\~1x        |\n| SP=4 (4 GPU)     | 32K        | 7,733    | **2.1x**    |\n| SP=4 (4 GPU)     | 64K        | 13,396   | **3.7x**    |\n\nAt the same 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 real benefit 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.\n\n> These results use only 4 GPUs with SP=4. With 8 GPUs (SP=8), you can push to even longer sequences — up to 256K+ tokens — or use 2D parallelism (SP=4, DP=2) to combine long-context training with data-parallel throughput.\n\n## Requirements\n\n- HF Accelerate: `deepspeed>=0.18.1 accelerate>=1.12`\n- HF Trainer: `deepspeed>=0.18.1 accelerate>=1.12 transformers>=5.0`\n- HF TRL: `deepspeed>=0.18.1 accelerate>=1.12 transformers>=5.0 trl>=0.18.0`\n\nUse `flash_attention_2` for Ampere GPUs, or `flash_attention_3` for Hopper GPUs. Wait for `flash_attention_4` on Blackwell 🕰.\n\n## Resources\n\n### Documentation\n\n- [Accelerate: Context Parallelism Guide](https://huggingface.co/docs/accelerate/concept_guides/context_parallelism)\n- [TRL: Distributing Training](https://huggingface.co/docs/trl/distributing_training)\n- [DeepSpeed Sequence Parallelism](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism)\n\n### Examples\n\n- [Accelerate ALST Example](https://github.com/huggingface/accelerate/tree/main/examples/alst_ulysses_sequence_parallelism)\n- [TRL Accelerate Configs](https://github.com/huggingface/trl/tree/main/examples/accelerate_configs)\n\n### Papers\n\n- [Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences](https://huggingface.co/papers/2506.13996)\n- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://huggingface.co/papers/2309.14509)\n\n### Related Blog Posts\n\n- [Accelerate ND-Parallel: A Guide to Efficient Multi-GPU Training](https://huggingface.co/blog/accelerate-nd-parallel)\n- [Understanding Ulysses and Ring Attention](https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention)\n- [Enabling Long-Context Training with Sequence Parallelism in Axolotl](https://huggingface.co/blog/axolotl-ai-co/long-context-with-sequence-parallelism-in-axolotl)\n",
  "authors": [],
  "captured_at": "2026-10-04T12:50:30.804479567+00:00",
  "category": "ai",
  "content_html": "<p>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.</p>\n<p>Ulysses Sequence Parallelism (part of the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/papers/2506.13996\">Arctic Long Sequence Training (ALST) protocol</a> 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.</p>\n<h2>Contents</h2>\n<ul>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#the-challenge-of-long-sequence-training\">The Challenge of Long Sequence Training</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#how-ulysses-works\">How Ulysses Works</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#integration-with-accelerate\">Integration with Accelerate</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#integration-with-transformers-trainer\">Integration with Transformers Trainer</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#integration-with-trl-sfttrainer\">Integration with TRL's SFTTrainer</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#comparing-ulysses-and-ring-attention\">Comparing Ulysses and Ring Attention</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#best-practices\">Best Practices</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#benchmarks\">Benchmarks</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/ulysses-sp#resources\">Resources</a></li>\n</ul>\n<h2>The Challenge of Long Sequence Training</h2>\n<p>The attention mechanism in transformers scales quadratically with sequence length. For a sequence of length n n, standard attention requires O(n2) O(n^2) FLOPs and O(n2) O(n^2) memory to compute and store the attention score matrix. Optimized implementations like <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/papers/2205.14135\">FlashAttention</a> reduce the <strong>memory</strong> to O(n) O(n) by tiling the computation and never materializing the full attention matrix—but the O(n2) O(n^2) compute remains. For very long sequences (32k+ tokens), even with FlashAttention, training still pushes the limits of single-GPU memory.</p>\n<p>Consider these scenarios where long-context training is essential:</p>\n<ul>\n<li><strong>Document understanding</strong>: Processing entire books, legal documents, or research papers</li>\n<li><strong>Code analysis</strong>: Understanding large codebases with multiple interconnected files</li>\n<li><strong>Reasoning tasks</strong>: Models that &quot;think&quot; step-by-step may generate thousands of tokens during inference</li>\n<li><strong>Retrieval-augmented generation</strong>: Incorporating many retrieved passages into the context</li>\n</ul>\n<p>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.</p>\n<h2>How Ulysses Works</h2>\n<p>Ulysses Sequence Parallelism (SP), introduced in the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/papers/2309.14509\">DeepSpeed Ulysses paper</a>, takes a clever approach: in addition to splitting on the sequence dimension, it also partitions the attention heads across GPUs.</p>\n<figure class=\"article-figure\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/ulysses_overview.png\" alt=\"Ulysses Sequence Parallelism Overview\" loading=\"eager\" decoding=\"async\" fetchpriority=\"high\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture><figcaption>Ulysses splits input sequences along the sequence dimension and uses all-to-all communication to exchange key-value pairs, enabling each GPU to compute a subset of attention heads. (<em><strong>Source: <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://www.snowflake.com/en/engineering-blog/ulysses-low-latency-llm-inference/\">Snowflake Engineering Blog</a></strong></em>)</figcaption></figure>\n<p>Here's how it works:</p>\n<ol>\n<li>\n<p><strong>Sequence Sharding</strong>: The input sequence is split along the sequence dimension across P P GPUs. Each GPU i i holds tokens [i⋅n/P,(i+1)⋅n/P) [i \\cdot n/P, (i+1) \\cdot n/P).</p>\n</li>\n<li>\n<p><strong>QKV Projection</strong>: Each GPU computes the query, key, and value projections for its local sequence chunk.</p>\n</li>\n<li>\n<p><strong>All-to-All Communication</strong>: An all-to-all collective operation redistributes the data so that each GPU holds <em>all</em> sequence positions after the projections, but only for a subset of attention heads.</p>\n</li>\n<li>\n<p><strong>Local Attention</strong>: Each GPU computes attention for its assigned heads using standard attention mechanisms (FlashAttention or SDPA).</p>\n</li>\n<li>\n<p><strong>All-to-All Communication</strong>: Another all-to-all operation reverses the redistribution, returning to sequence-sharded format.</p>\n</li>\n<li>\n<p><strong>Output Projection</strong>: Each GPU computes the output projection for its local sequence chunk.</p>\n</li>\n</ol>\n<p>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.</p>\n<h3>Communication Complexity</h3>\n<p>Ulysses requires two all-to-all operations per attention layer, with total communication volume of O(n⋅d/P) O(n \\cdot d / P) per GPU, where:</p>\n<ul>\n<li>n n is the sequence length</li>\n<li>d d is the hidden dimension</li>\n<li>P P is the parallelism degree</li>\n</ul>\n<p>Ring Attention communicates O(n⋅d) O(n \\cdot d) per GPU — a factor of P P more — via P−1 P-1 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 P−1 P-1 hops.</p>\n<h2>Integration with Accelerate</h2>\n<p>Accelerate provides the foundation for Ulysses sequence parallelism through its <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/accelerate/concept_guides/sequence_parallelism#alstulysses-sp-backend-configuration\"><code>ParallelismConfig</code></a> class and DeepSpeed integration.</p>\n<h3>Configuration</h3>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">accelerate</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">Accelerator</span></span>\n<span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">accelerate</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-import-name syntax-python\">utils</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">ParallelismConfig</span><span class=\"syntax-punctuation syntax-separator syntax-import-list syntax-python\">,</span> <span class=\"syntax-meta syntax-generic-name syntax-python\">DeepSpeedSequenceParallelConfig</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">ParallelismConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_backend</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">deepspeed<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_size</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">4</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Split across 4 GPUs\n</span>    <span class=\"syntax-variable syntax-parameter syntax-python\">dp_shard_size</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Must satisfy: dp_replicate × dp_shard × sp_size = num_processes\n</span>    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_handler</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">DeepSpeedSequenceParallelConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n        <span class=\"syntax-variable syntax-parameter syntax-python\">sp_seq_length</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">None</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> None for variable-length sequences\n</span>        <span class=\"syntax-variable syntax-parameter syntax-python\">sp_seq_length_is_variable</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n        <span class=\"syntax-variable syntax-parameter syntax-python\">sp_attn_implementation</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">flash_attention_2<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> or &quot;sdpa&quot;\n</span>    </span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">accelerator</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Accelerator</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-variable syntax-parameter syntax-python\">parallelism_config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<h3>Key Parameters</h3>\n<table>\n<thead>\n<tr>\n<th>Parameter</th>\n<th>Description</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><code>sp_size</code></td>\n<td>Number of GPUs for sequence parallelism</td>\n</tr>\n<tr>\n<td><code>sp_backend</code></td>\n<td>Must be <code>&quot;deepspeed&quot;</code> for Ulysses</td>\n</tr>\n<tr>\n<td><code>sp_seq_length_is_variable</code></td>\n<td>Set to <code>True</code> for varying sequence lengths across batches</td>\n</tr>\n<tr>\n<td><code>sp_attn_implementation</code></td>\n<td><code>&quot;flash_attention_2&quot;</code>, <code>&quot;flash_attention_3&quot;</code>, or <code>&quot;sdpa&quot;</code></td>\n</tr>\n</tbody>\n</table>\n<h3>Using the Accelerator</h3>\n<p>When you call <code>accelerator.prepare()</code>, Ulysses is automatically set up:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">transformers</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">AutoModelForCausalLM</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">AutoModelForCausalLM</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">from_pretrained</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">meta-llama/Llama-3.1-8B<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">optimizer</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">optim</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">AdamW</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">parameters</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">lr</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-float syntax-python\">1e-5</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> This registers the model with Ulysses and wraps the dataloader\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span></span>, <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">optimizer</span></span>, <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dataloader</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">accelerator</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">prepare</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">optimizer</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dataloader</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>The <code>prepare()</code> call:</p>\n<ol>\n<li>Registers the model with DeepSpeed's <code>UlyssesSPAttentionHF</code></li>\n<li>Wraps the dataloader with <code>UlyssesSPDataLoaderAdapter</code> to handle sequence sharding</li>\n<li>Automatically injects <code>shift_labels</code> for correct loss computation</li>\n</ol>\n<h3>Loss Aggregation</h3>\n<p>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 <code>Trainer</code> or TRL's <code>SFTTrainer</code>, this is handled automatically—the code below is only needed when writing a custom Accelerate training loop:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_size</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_size</span></span>\n<span class=\"syntax-meta syntax-statement syntax-if syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-conditional syntax-python\">if</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_size</span></span> <span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-section syntax-block syntax-conditional syntax-python\">:</span></span>\n    <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">deepspeed</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-import-name syntax-python\">utils</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">groups</span></span>\n\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_group</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">groups</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">_get_sequence_parallel_group</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n    <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Gather losses and token counts from all SP ranks\n</span>    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">losses_per_rank</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">distributed</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">nn</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">functional</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">all_gather</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">loss</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">group</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_group</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">good_tokens</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-group syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-group syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">batch</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">shift_labels<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span> <span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">!=</span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">100</span><span class=\"syntax-punctuation syntax-section syntax-group syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">view</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">sum</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">good_tokens_per_rank</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">distributed</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">nn</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">functional</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">all_gather</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">good_tokens</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">group</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_group</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n    <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Weighted aggregation\n</span>    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">total_loss</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">sum</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n        <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">losses_per_rank</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">i</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">*</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">good_tokens_per_rank</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">i</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span>\n        <span class=\"syntax-meta syntax-expression syntax-generator syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-for syntax-generator syntax-python\">for</span> <span class=\"syntax-meta syntax-generic-name syntax-python\">i</span> <span class=\"syntax-keyword syntax-control syntax-flow syntax-for syntax-in syntax-python\">in</span></span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">range</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_size</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-keyword syntax-control syntax-flow syntax-if syntax-inline syntax-python\">if</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">good_tokens_per_rank</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">i</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span> <span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">0</span>\n    </span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">loss</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">total_loss</span></span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">/</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">max</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">sum</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">good_tokens_per_rank</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">accelerator</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">backward</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">loss</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<blockquote>\n<p>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).</p>\n</blockquote>\n<blockquote>\n<p>Both Ulysses and Ring Attention use <code>position_ids</code> instead of <code>attention_mask</code> 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 O(n) O(n) memory instead of O(n2) O(n^2). During evaluation/inference, DeepSpeed's SP attention layer can bypass the SP operations entirely (via <code>disable_in_eval</code>) and fall back to the model's default attention implementation.</p>\n</blockquote>\n<h2>Integration with Transformers Trainer</h2>\n<p>The Transformers <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/transformers/main_classes/trainer\"><code>Trainer</code></a> provides seamless Ulysses integration through <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.parallelism_config\"><code>TrainingArguments.parallelism_config</code></a>. 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.</p>\n<h3>Configuration</h3>\n<p>Just pass the same <code>parallelism_config</code> from above into <code>TrainingArguments</code>:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">transformers</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">TrainingArguments</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">training_args</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">TrainingArguments</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">parallelism_config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> same ParallelismConfig as above\n</span>    <span class=\"syntax-variable syntax-parameter syntax-python\">per_device_train_batch_size</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<h3>What the Trainer Handles Automatically</h3>\n<ol>\n<li>\n<p><strong>Dataloader Wrapping</strong>: After model preparation, the Trainer wraps the dataloader with <code>UlyssesSPDataLoaderAdapter</code></p>\n</li>\n<li>\n<p><strong>Loss Computation</strong>: The <code>compute_loss</code> method detects SP mode and routes to specialized <code>_deepspeed_sp_compute_loss</code> which handles:</p>\n<ul>\n<li>Gathering losses across SP ranks</li>\n<li>Computing valid token counts per rank</li>\n<li>Weighted loss aggregation</li>\n</ul>\n</li>\n<li>\n<p><strong>Batch Size Calculation</strong>: The effective data parallel world size accounts for SP:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dp_world_size</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">world_size</span></span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">/</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">/</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sp_size</span></span>\n</span></span></code></pre>\n</li>\n<li>\n<p><strong>Dataloader Length Adjustment</strong>: Training step calculations are adjusted for SP's effect on iteration count</p>\n</li>\n</ol>\n<h3>Launch Command</h3>\n<p>Use an accelerate config file or command-line arguments:</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">accelerate</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> launch <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>config_file</span> deepspeed_ulysses.yaml <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span>    train.py <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>per_device_train_batch_size</span> 1</span>\n</span></span></code></pre>\n<h2>Integration with TRL SFTTrainer</h2>\n<p>TRL's <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/trl/sft_trainer\"><code>SFTTrainer</code></a> builds on the Transformers Trainer and adds specific optimizations for supervised fine-tuning with long sequences.</p>\n<h3>Configuration</h3>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">trl</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">SFTConfig</span><span class=\"syntax-punctuation syntax-separator syntax-import-list syntax-python\">,</span> <span class=\"syntax-meta syntax-generic-name syntax-python\">SFTTrainer</span></span>\n<span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">accelerate</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-import-name syntax-python\">utils</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">ParallelismConfig</span><span class=\"syntax-punctuation syntax-separator syntax-import-list syntax-python\">,</span> <span class=\"syntax-meta syntax-generic-name syntax-python\">DeepSpeedSequenceParallelConfig</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">ParallelismConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_backend</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">deepspeed<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_size</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">dp_shard_size</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> 2D parallelism: SP × DP = 4 GPUs\n</span>    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_handler</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">DeepSpeedSequenceParallelConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n        <span class=\"syntax-variable syntax-parameter syntax-python\">sp_seq_length_is_variable</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n        <span class=\"syntax-variable syntax-parameter syntax-python\">sp_attn_implementation</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">flash_attention_2<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    </span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">training_args</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">SFTConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-constant syntax-language syntax-python\">...</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">parallelism_config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">max_length</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">32768</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">pad_to_multiple_of</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Must equal sp_size\n</span>    <span class=\"syntax-variable syntax-parameter syntax-python\">per_device_train_batch_size</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">trainer</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">SFTTrainer</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">model</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">args</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">training_args</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">train_dataset</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dataset</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">trainer</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">train</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<h3>Key SFTConfig Parameters for Ulysses</h3>\n<table>\n<thead>\n<tr>\n<th>Parameter</th>\n<th>Description</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><code>pad_to_multiple_of</code></td>\n<td>Must equal <code>sp_size</code> to ensure sequence divisibility</td>\n</tr>\n<tr>\n<td><code>max_length</code></td>\n<td>Global sequence length (before splitting across GPUs)</td>\n</tr>\n<tr>\n<td><code>packing</code></td>\n<td>Works well with SP — packing reduces padding waste, especially for variable-length sequences</td>\n</tr>\n</tbody>\n</table>\n<h3>Accelerate Config File</h3>\n<p>Create <code>alst_ulysses_4gpu.yaml</code>:</p>\n<pre><code class=\"language-yaml\"><span class=\"code-snippet\" data-language=\"YAML\"><span class=\"syntax-source syntax-yaml\"><span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">compute_environment</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">LOCAL_MACHINE</span>\n<span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">distributed_type</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">DEEPSPEED</span>\n<span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">mixed_precision</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">bf16</span>\n<span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">num_processes</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-yaml\">4</span>\n<span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">deepspeed_config</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">zero_stage</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-yaml\">3</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">seq_parallel_communication_data_type</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">bf16</span>\n<span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">parallelism_config</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">parallelism_config_sp_size</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-yaml\">2</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">parallelism_config_sp_backend</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">deepspeed</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">parallelism_config_dp_shard_size</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-yaml\">2</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">parallelism_config_sp_seq_length_is_variable</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-constant syntax-language syntax-boolean syntax-yaml\">true</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">parallelism_config_sp_attn_implementation</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">flash_attention_2</span>\n</span></span></code></pre>\n<h3>Complete Training Command</h3>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">accelerate</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> launch<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>config_file</span> alst_ulysses_4gpu.yaml <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span>    trl/scripts/sft.py <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>model_name_or_path</span> meta-llama/Llama-3.1-8B <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>dataset_name</span> trl-lib/Capybara <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>max_length</span> 32768 <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>packing</span> <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>pad_to_multiple_of</span> 2 <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\">    --</span>per_device_train_batch_size</span> 1</span>\n</span></span></code></pre>\n<h3>Shift Labels Handling</h3>\n<p>The SFTTrainer automatically handles pre-shifted labels when Ulysses is enabled:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> When using SP, labels are pre-shifted by the dataloader adapter\n</span><span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> The trainer detects this and uses shift_labels directly\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">labels</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">inputs</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">labels<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span> <span class=\"syntax-meta syntax-statement syntax-if syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-conditional syntax-python\">if</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">shift_labels<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span> <span class=\"syntax-keyword syntax-operator syntax-logical syntax-python\">not</span> <span class=\"syntax-keyword syntax-operator syntax-logical syntax-python\">in</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">inputs</span></span> <span class=\"syntax-keyword syntax-control syntax-flow syntax-else syntax-inline syntax-python\">else</span> <span class=\"syntax-constant syntax-language syntax-python\">None</span></span>\n\n<span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Loss computation uses the pre-shifted labels\n</span><span class=\"syntax-meta syntax-statement syntax-if syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-conditional syntax-python\">if</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">shift_labels<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span> <span class=\"syntax-keyword syntax-operator syntax-logical syntax-python\">in</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">inputs</span></span><span class=\"syntax-punctuation syntax-section syntax-block syntax-conditional syntax-python\">:</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">shift_logits</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">outputs</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">logits</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">contiguous</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">shift_labels</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">inputs</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">shift_labels<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span>\n<span class=\"syntax-meta syntax-statement syntax-conditional syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-conditional syntax-python\">else</span><span class=\"syntax-punctuation syntax-section syntax-block syntax-python\">:</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">shift_logits</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">outputs</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">logits</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-constant syntax-language syntax-python\">...</span>, <span class=\"syntax-punctuation syntax-separator syntax-slice syntax-python\">:</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span>, <span class=\"syntax-punctuation syntax-separator syntax-slice syntax-python\">:</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">contiguous</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">shift_labels</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">labels</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-constant syntax-language syntax-python\">...</span>, <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-slice syntax-python\">:</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">contiguous</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<h2>Comparing Ulysses and Ring Attention</h2>\n<p>Both Ulysses and Ring Attention enable long-context training, but they have different characteristics:</p>\n<table>\n<thead>\n<tr>\n<th>Aspect</th>\n<th>Ulysses (DeepSpeed)</th>\n<th>Ring Attention (FSDP2)</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><strong>Parallelism Method</strong></td>\n<td>Attention head partitioning</td>\n<td>Ring-based KV exchange</td>\n</tr>\n<tr>\n<td><strong>Backend</strong></td>\n<td>DeepSpeed ZeRO</td>\n<td>PyTorch FSDP2</td>\n</tr>\n<tr>\n<td><strong>Attention Support</strong></td>\n<td>FlashAttention 2/3, SDPA</td>\n<td>SDPA only</td>\n</tr>\n<tr>\n<td><strong>Communication</strong></td>\n<td>Two <code>all-to-all</code>s per layer</td>\n<td>P2P ring communication</td>\n</tr>\n<tr>\n<td><strong>Comm volume per GPU</strong></td>\n<td>O(total_seq x hidden / sp_size)</td>\n<td>O(total_seq x hidden)</td>\n</tr>\n<tr>\n<td><strong>Sequence Divisibility</strong></td>\n<td><code>sp_size</code></td>\n<td><code>cp_size * 2</code></td>\n</tr>\n<tr>\n<td><strong>Num Head Constraint</strong></td>\n<td><code>num_heads &gt;= sp_size</code></td>\n<td>None</td>\n</tr>\n</tbody>\n</table>\n<h3>When to Choose Ulysses vs Ring Attention</h3>\n<p>Since switching between the two only requires changing the accelerate config, we recommend trying both and comparing performance and memory usage on your specific setup. The main constraint is that Ulysses requires <code>num_heads &gt;= sp_size</code>, while Ring Attention has no such limitation.</p>\n<h2>Best Practices</h2>\n<h3>1. Sequence Length Divisibility</h3>\n<p>Always ensure your sequence length is divisible by <code>sp_size</code>:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">training_args</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">SFTConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">pad_to_multiple_of</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">4</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> For sp_size=4\n</span>    <span class=\"syntax-variable syntax-parameter syntax-python\">max_length</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">32768</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Must be divisible by 4\n</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<h3>2. Use Flash Attention</h3>\n<p>Flash Attention 2 provides cleaner output and better performance than SDPA:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">parallelism_config</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">ParallelismConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">sp_handler</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">DeepSpeedSequenceParallelConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n        <span class=\"syntax-variable syntax-parameter syntax-python\">sp_attn_implementation</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">flash_attention_2<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    </span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>Use Flash Attention 3 for Hopper and look out for Flash Attention 4 release for Blackwell (FA2 on Blackwell is quite slow).</p>\n<h3>3. Combine with DeepSpeed ZeRO</h3>\n<p>For very large models, combine Ulysses with ZeRO Stage 3:</p>\n<pre><code class=\"language-yaml\"><span class=\"code-snippet\" data-language=\"YAML\"><span class=\"syntax-source syntax-yaml\"><span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">deepspeed_config</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">zero_stage</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-yaml\">3</span>\n  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">offload_optimizer</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span>\n    <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">device</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">cpu</span>\n</span></span></code></pre>\n<p>If the model is huge, you can offload the params as well by adding to the above:</p>\n<pre><code class=\"language-yaml\"><span class=\"code-snippet\" data-language=\"YAML\"><span class=\"syntax-source syntax-yaml\">  <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">offload_param</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span>\n    <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\"><span class=\"syntax-entity syntax-name syntax-tag syntax-yaml\">device</span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-mapping syntax-yaml\">:</span> <span class=\"syntax-string syntax-unquoted syntax-plain syntax-out syntax-yaml\">cpu</span>\n</span></span></code></pre>\n<h3>5. Use memory fragmentation-friendly PyTorch allocator</h3>\n<p>This environment variable will allow for a longer sequence length:</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-storage syntax-modifier syntax-shell\">export</span> <span class=\"syntax-variable syntax-other syntax-readwrite syntax-assignment syntax-shell\">PYTORCH_ALLOC_CONF</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-shell\">=</span><span class=\"syntax-string syntax-unquoted syntax-shell\">expandable_segments:True</span></span>\n</span></span></code></pre>\n<h3>6. 2D Parallelism Configuration</h3>\n<p>Balance SP and DP for your GPU count:</p>\n<table>\n<thead>\n<tr>\n<th>GPUs</th>\n<th><code>sp_size</code></th>\n<th><code>dp_shard_size</code></th>\n<th>Use Case</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>4</td>\n<td>2</td>\n<td>2</td>\n<td>Balanced throughput and sequence length</td>\n</tr>\n<tr>\n<td>4</td>\n<td>4</td>\n<td>1</td>\n<td>Maximum sequence length</td>\n</tr>\n<tr>\n<td>8</td>\n<td>2</td>\n<td>4</td>\n<td>Higher throughput with moderate sequence length</td>\n</tr>\n<tr>\n<td>8</td>\n<td>4</td>\n<td>2</td>\n<td>Longer sequences with moderate throughput</td>\n</tr>\n</tbody>\n</table>\n<p>Remember: <code>dp_replicate_size × dp_shard_size × sp_size = num_processes</code></p>\n<h3>7. Liger-Kernel</h3>\n<p>If your desired model architecture is supported by <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/linkedin/Liger-Kernel\">Liger-Kernel</a>, it is fully compatible with Ulysses SP and can be enabled with a single flag:</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">training_args</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">SFTConfig</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-variable syntax-parameter syntax-python\">use_liger_kernel</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>The main memory saving comes from <code>FusedLinearCrossEntropy</code> which avoids materializing the full logits tensor during loss calculation. The savings grow with longer sequences where the logits tensor is larger.</p>\n<p>Additionally, you can enable <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/#tiled-mlp-computation\"><code>TiledMLP</code></a> to further extend sequence length — like <code>FusedLinearCrossEntropy</code>, it saves working memory by tiling large matrix operations.</p>\n<h3>8. Token Distribution Across Ranks</h3>\n<p>You don't need to 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 reasonably sized dataset, the distribution evens out statistically over training.</p>\n<h2>Benchmarks</h2>\n<p>To quantify the benefits of Ulysses SP, we trained <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/Qwen/Qwen3-4B\">Qwen3-4B</a> on the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/sedthh/gutenberg_english\">Gutenberg English</a> 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 attention backend.</p>\n<h3>Setup</h3>\n<table>\n<thead>\n<tr>\n<th>Config</th>\n<th>GPUs</th>\n<th>SP</th>\n<th>DP</th>\n<th>Seq Length</th>\n<th>Grad Acc</th>\n<th>Global Batch</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Baseline</td>\n<td>1</td>\n<td>1</td>\n<td>1</td>\n<td>8K</td>\n<td>8</td>\n<td>8</td>\n</tr>\n<tr>\n<td>SP=4</td>\n<td>4</td>\n<td>4</td>\n<td>1</td>\n<td>8K</td>\n<td>8</td>\n<td>8</td>\n</tr>\n<tr>\n<td>SP=4</td>\n<td>4</td>\n<td>4</td>\n<td>1</td>\n<td>32K</td>\n<td>8</td>\n<td>8</td>\n</tr>\n<tr>\n<td>SP=4</td>\n<td>4</td>\n<td>4</td>\n<td>1</td>\n<td>64K</td>\n<td>8</td>\n<td>8</td>\n</tr>\n<tr>\n<td>SP=4</td>\n<td>4</td>\n<td>4</td>\n<td>1</td>\n<td>96K</td>\n<td>8</td>\n<td>8</td>\n</tr>\n</tbody>\n</table>\n<p>The benchmark runs in the table above use the same global batch size (8 micro-batches), cosine learning-rate schedule, and seed, so those benchmark loss curves are directly comparable.</p>\n<h3>Loss Curve Matching Diagnostics (4 GPU)</h3>\n<p>To verify SP-vs-DP loss equivalence, we ran controlled 4-GPU A/B experiments with identical seed, model, optimizer, learning-rate schedule, and data order.</p>\n<h4>Methodology for Fair DP vs SP Comparison</h4>\n<p>Compared setups:</p>\n<ul>\n<li><code>DP=4, SP=1, GAS=1</code> (baseline)</li>\n<li><code>DP=1, SP=4, GAS=4</code> (Ulysses SP)</li>\n</ul>\n<p>For fair comparison, <code>GAS</code> must scale with <code>SP</code>:</p>\n<ul>\n<li>Ulysses SP splits the sequence across <code>SP</code> ranks, so each SP rank sees roughly <code>1/SP</code> of the sequence tokens per micro-step.</li>\n<li>If <code>GAS</code> is unchanged, each optimizer step in SP aggregates fewer total tokens than the DP baseline.</li>\n<li>Setting <code>GAS=SP</code> keeps effective tokens per optimizer step matched:\n<ul>\n<li>DP tokens/step: <code>dp_world_size * micro_batch * seq_len * GAS = 4 * B * L * 1</code></li>\n<li>SP tokens/step: <code>dp_world_size * micro_batch * (L/SP) * GAS * SP_ranks = 1 * B * (L/4) * 4 * 4 = 4 * B * L</code></li>\n</ul>\n</li>\n</ul>\n<figure class=\"article-figure\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/gutenberg_canonical_loss_dp4_vs_sp4.png\" alt=\"Canonical loss on Gutenberg for DP=4 vs SP=4\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture><figcaption>On Gutenberg text (20 steps), canonical loss matches within logging precision between <code>DP=4,SP=1,GAS=1</code> and <code>DP=1,SP=4,GAS=4</code>.</figcaption></figure>\n<p>Measured over 20 steps on 4 GPUs in controlled equivalence harnesses:</p>\n<table>\n<thead>\n<tr>\n<th>Harness</th>\n<th>Metric</th>\n<th>DP vs SP setting</th>\n<th>Mean abs diff</th>\n<th>Max abs diff</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><code>Trainer</code></td>\n<td><code>loss</code></td>\n<td>DP=4, SP=1 vs DP=1, SP=4</td>\n<td>0.0054</td>\n<td>0.0131</td>\n</tr>\n<tr>\n<td><code>SFTTrainer</code></td>\n<td>logged <code>loss</code></td>\n<td>DP=4, SP=1 vs DP=1, SP=4</td>\n<td>0.0811</td>\n<td>0.0812</td>\n</tr>\n<tr>\n<td><code>SFTTrainer</code></td>\n<td>canonical NLL</td>\n<td>DP=4, SP=1 vs DP=1, SP=4</td>\n<td>0.000004</td>\n<td>0.000005</td>\n</tr>\n</tbody>\n</table>\n<p>Takeaway: under matched token budget, SP and non-SP match on canonical token-normalized loss. The remaining difference is in trainer-reported logging (<code>loss</code>), not in the underlying cross-entropy objective.</p>\n<h3>Memory Reduction</h3>\n<figure class=\"article-figure\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/memory_usage.png\" alt=\"Peak GPU Memory per Rank\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture><figcaption>SP=4 reduces per-GPU memory by 3.3x at the same sequence length, enabling training at up to 96K tokens on 4× H100 80GB. At 128K, the model OOMs.</figcaption></figure>\n<table>\n<thead>\n<tr>\n<th>Config</th>\n<th>Seq Length</th>\n<th>Peak Memory</th>\n<th>Notes</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>DP=4 (4 GPU)</td>\n<td>8K</td>\n<td>22.4 GB</td>\n<td>Baseline — no SP</td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>8K</td>\n<td>22.8 GB</td>\n<td>Similar memory at same seq length</td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>32K</td>\n<td>35.0 GB</td>\n<td>4x longer than DP baseline</td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>64K</td>\n<td>50.5 GB</td>\n<td>8x longer than DP baseline</td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>96K</td>\n<td>66.0 GB</td>\n<td><strong>12x longer than DP baseline</strong></td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>128K</td>\n<td>OOM</td>\n<td>Exceeds 80 GB limit</td>\n</tr>\n</tbody>\n</table>\n<p>At 8K tokens, DP=4 and SP=4 use nearly the same memory per GPU (~22 GB with ZeRO-3). The advantage of SP is that it enables scaling to much longer sequences: at 96K tokens (12x longer), peak memory is 66 GB — still within the H100's 80 GB capacity. At 128K, the model OOMs, establishing the practical limit for this configuration. DP=4 without SP cannot scale beyond 8K for this model.</p>\n<h3>Throughput</h3>\n<figure class=\"article-figure\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/throughput.png\" alt=\"Training Throughput\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture><figcaption>Longer sequences with SP process dramatically more tokens per second. SP=4 at 64K achieves 3.7x the throughput of the baseline.</figcaption></figure>\n<table>\n<thead>\n<tr>\n<th>Config</th>\n<th>Seq Length</th>\n<th>Tokens/s</th>\n<th>vs Baseline</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Baseline (1 GPU)</td>\n<td>8K</td>\n<td>3,633</td>\n<td>—</td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>8K</td>\n<td>3,933</td>\n<td>~1x</td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>32K</td>\n<td>7,733</td>\n<td><strong>2.1x</strong></td>\n</tr>\n<tr>\n<td>SP=4 (4 GPU)</td>\n<td>64K</td>\n<td>13,396</td>\n<td><strong>3.7x</strong></td>\n</tr>\n</tbody>\n</table>\n<p>At the same 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 real benefit 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.</p>\n<blockquote>\n<p>These results use only 4 GPUs with SP=4. With 8 GPUs (SP=8), you can push to even longer sequences — up to 256K+ tokens — or use 2D parallelism (SP=4, DP=2) to combine long-context training with data-parallel throughput.</p>\n</blockquote>\n<h2>Requirements</h2>\n<ul>\n<li>HF Accelerate: <code>deepspeed&gt;=0.18.1 accelerate&gt;=1.12</code></li>\n<li>HF Trainer: <code>deepspeed&gt;=0.18.1 accelerate&gt;=1.12 transformers&gt;=5.0</code></li>\n<li>HF TRL: <code>deepspeed&gt;=0.18.1 accelerate&gt;=1.12 transformers&gt;=5.0 trl&gt;=0.18.0</code></li>\n</ul>\n<p>Use <code>flash_attention_2</code> for Ampere GPUs, or <code>flash_attention_3</code> for Hopper GPUs. Wait for <code>flash_attention_4</code> on Blackwell 🕰.</p>\n<h2>Resources</h2>\n<h3>Documentation</h3>\n<ul>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/accelerate/concept_guides/context_parallelism\">Accelerate: Context Parallelism Guide</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/trl/distributing_training\">TRL: Distributing Training</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism\">DeepSpeed Sequence Parallelism</a></li>\n</ul>\n<h3>Examples</h3>\n<ul>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/huggingface/accelerate/tree/main/examples/alst_ulysses_sequence_parallelism\">Accelerate ALST Example</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/huggingface/trl/tree/main/examples/accelerate_configs\">TRL Accelerate Configs</a></li>\n</ul>\n<h3>Papers</h3>\n<ul>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/papers/2506.13996\">Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/papers/2309.14509\">DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models</a></li>\n</ul>\n<h3>Related Blog Posts</h3>\n<ul>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/accelerate-nd-parallel\">Accelerate ND-Parallel: A Guide to Efficient Multi-GPU Training</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention\">Understanding Ulysses and Ring Attention</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/axolotl-ai-co/long-context-with-sequence-parallelism-in-axolotl\">Enabling Long-Context Training with Sequence Parallelism in Axolotl</a></li>\n</ul>\n",
  "content_markdown": "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.\n\nUlysses Sequence Parallelism (part of the [Arctic Long Sequence Training (ALST) protocol](https://huggingface.co/papers/2506.13996) 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.\n\n## Contents\n\n- [The Challenge of Long Sequence Training](https://huggingface.co/blog/ulysses-sp#the-challenge-of-long-sequence-training)\n- [How Ulysses Works](https://huggingface.co/blog/ulysses-sp#how-ulysses-works)\n- [Integration with Accelerate](https://huggingface.co/blog/ulysses-sp#integration-with-accelerate)\n- [Integration with Transformers Trainer](https://huggingface.co/blog/ulysses-sp#integration-with-transformers-trainer)\n- [Integration with TRL's SFTTrainer](https://huggingface.co/blog/ulysses-sp#integration-with-trl-sfttrainer)\n- [Comparing Ulysses and Ring Attention](https://huggingface.co/blog/ulysses-sp#comparing-ulysses-and-ring-attention)\n- [Best Practices](https://huggingface.co/blog/ulysses-sp#best-practices)\n- [Benchmarks](https://huggingface.co/blog/ulysses-sp#benchmarks)\n- [Resources](https://huggingface.co/blog/ulysses-sp#resources)\n\n## The Challenge of Long Sequence Training\n\nThe attention mechanism in transformers scales quadratically with sequence length. For a sequence of length n n, standard attention requires O(n2) O(n^2) FLOPs and O(n2) O(n^2) memory to compute and store the attention score matrix. Optimized implementations like [FlashAttention](https://huggingface.co/papers/2205.14135) reduce the **memory** to O(n) O(n) by tiling the computation and never materializing the full attention matrix—but the O(n2) O(n^2) compute remains. For very long sequences (32k+ tokens), even with FlashAttention, training still pushes the limits of single-GPU memory.\n\nConsider these scenarios where long-context training is essential:\n\n- **Document understanding**: Processing entire books, legal documents, or research papers\n- **Code analysis**: Understanding large codebases with multiple interconnected files\n- **Reasoning tasks**: Models that \"think\" step-by-step may generate thousands of tokens during inference\n- **Retrieval-augmented generation**: Incorporating many retrieved passages into the context\n\nTraditional 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.\n\n## How Ulysses Works\n\nUlysses Sequence Parallelism (SP), introduced in the [DeepSpeed Ulysses paper](https://huggingface.co/papers/2309.14509), takes a clever approach: in addition to splitting on the sequence dimension, it also partitions the attention heads across GPUs.\n\n![Ulysses Sequence Parallelism Overview](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/ulysses_overview.png)\\\nUlysses splits input sequences along the sequence dimension and uses all-to-all communication to exchange key-value pairs, enabling each GPU to compute a subset of attention heads. (***Source: [Snowflake Engineering Blog](https://www.snowflake.com/en/engineering-blog/ulysses-low-latency-llm-inference/)***)\n\nHere's how it works:\n\n1. **Sequence Sharding**: The input sequence is split along the sequence dimension across P P GPUs. Each GPU i i holds tokens \\[i⋅n/P,(i+1)⋅n/P) \\[i \\\\cdot n/P, (i+1) \\\\cdot n/P).\n\n2. **QKV Projection**: Each GPU computes the query, key, and value projections for its local sequence chunk.\n\n3. **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.\n\n4. **Local Attention**: Each GPU computes attention for its assigned heads using standard attention mechanisms (FlashAttention or SDPA).\n\n5. **All-to-All Communication**: Another all-to-all operation reverses the redistribution, returning to sequence-sharded format.\n\n6. **Output Projection**: Each GPU computes the output projection for its local sequence chunk.\n\nThe 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.\n\n### Communication Complexity\n\nUlysses requires two all-to-all operations per attention layer, with total communication volume of O(n⋅d/P) O(n \\\\cdot d / P) per GPU, where:\n\n- n n is the sequence length\n- d d is the hidden dimension\n- P P is the parallelism degree\n\nRing Attention communicates O(n⋅d) O(n \\\\cdot d) per GPU — a factor of P P more — via P−1 P-1 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 P−1 P-1 hops.\n\n## Integration with Accelerate\n\nAccelerate provides the foundation for Ulysses sequence parallelism through its [`ParallelismConfig`](https://huggingface.co/docs/accelerate/concept_guides/sequence_parallelism#alstulysses-sp-backend-configuration) class and DeepSpeed integration.\n\n### Configuration\n\n```python\nfrom accelerate import Accelerator\nfrom accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig\n\nparallelism_config = ParallelismConfig(\n    sp_backend=\"deepspeed\",\n    sp_size=4,  # Split across 4 GPUs\n    dp_shard_size=1,  # Must satisfy: dp_replicate × dp_shard × sp_size = num_processes\n    sp_handler=DeepSpeedSequenceParallelConfig(\n        sp_seq_length=None,  # None for variable-length sequences\n        sp_seq_length_is_variable=True,\n        sp_attn_implementation=\"flash_attention_2\",  # or \"sdpa\"\n    ),\n)\n\naccelerator = Accelerator(parallelism_config=parallelism_config)\n```\n\n### Key Parameters\n\n| Parameter                   | Description                                               |\n| --------------------------- | --------------------------------------------------------- |\n| `sp_size`                   | Number of GPUs for sequence parallelism                   |\n| `sp_backend`                | Must be `\"deepspeed\"` for Ulysses                         |\n| `sp_seq_length_is_variable` | Set to `True` for varying sequence lengths across batches |\n| `sp_attn_implementation`    | `\"flash_attention_2\"`, `\"flash_attention_3\"`, or `\"sdpa\"` |\n\n### Using the Accelerator\n\nWhen you call `accelerator.prepare()`, Ulysses is automatically set up:\n\n```python\nfrom transformers import AutoModelForCausalLM\n\nmodel = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-3.1-8B\")\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)\n\n# This registers the model with Ulysses and wraps the dataloader\nmodel, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)\n```\n\nThe `prepare()` call:\n\n1. Registers the model with DeepSpeed's `UlyssesSPAttentionHF`\n2. Wraps the dataloader with `UlyssesSPDataLoaderAdapter` to handle sequence sharding\n3. Automatically injects `shift_labels` for correct loss computation\n\n### Loss Aggregation\n\nWith 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:\n\n```python\nsp_size = parallelism_config.sp_size\nif sp_size > 1:\n    from deepspeed.utils import groups\n\n    sp_group = groups._get_sequence_parallel_group()\n\n    # Gather losses and token counts from all SP ranks\n    losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)\n    good_tokens = (batch[\"shift_labels\"] != -100).view(-1).sum()\n    good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)\n\n    # Weighted aggregation\n    total_loss = sum(\n        losses_per_rank[i] * good_tokens_per_rank[i]\n        for i in range(sp_size)\n        if good_tokens_per_rank[i] > 0\n    )\n    loss = total_loss / max(sum(good_tokens_per_rank), 1)\n\naccelerator.backward(loss)\n```\n\n> 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).\n\n> Both Ulysses and Ring Attention use `position_ids` instead of `attention_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 O(n) O(n) memory instead of O(n2) O(n^2). During evaluation/inference, DeepSpeed's SP attention layer can bypass the SP operations entirely (via `disable_in_eval`) and fall back to the model's default attention implementation.\n\n## Integration with Transformers Trainer\n\nThe Transformers [`Trainer`](https://huggingface.co/docs/transformers/main_classes/trainer) provides seamless Ulysses integration through [`TrainingArguments.parallelism_config`](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.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.\n\n### Configuration\n\nJust pass the same `parallelism_config` from above into `TrainingArguments`:\n\n```python\nfrom transformers import TrainingArguments\n\ntraining_args = TrainingArguments(\n    parallelism_config=parallelism_config,  # same ParallelismConfig as above\n    per_device_train_batch_size=1,\n)\n```\n\n### What the Trainer Handles Automatically\n\n1. **Dataloader Wrapping**: After model preparation, the Trainer wraps the dataloader with `UlyssesSPDataLoaderAdapter`\n\n2. **Loss Computation**: The `compute_loss` method detects SP mode and routes to specialized `_deepspeed_sp_compute_loss` which handles:\n\n   - Gathering losses across SP ranks\n   - Computing valid token counts per rank\n   - Weighted loss aggregation\n3. **Batch Size Calculation**: The effective data parallel world size accounts for SP:\n\n   ```python\n   dp_world_size = world_size // sp_size\n   ```\n\n4. **Dataloader Length Adjustment**: Training step calculations are adjusted for SP's effect on iteration count\n\n### Launch Command\n\nUse an accelerate config file or command-line arguments:\n\n```bash\naccelerate launch \\\n    --config_file deepspeed_ulysses.yaml \\\n    train.py \\\n    --per_device_train_batch_size 1\n```\n\n## Integration with TRL SFTTrainer\n\nTRL's [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) builds on the Transformers Trainer and adds specific optimizations for supervised fine-tuning with long sequences.\n\n### Configuration\n\n```python\nfrom trl import SFTConfig, SFTTrainer\nfrom accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig\n\nparallelism_config = ParallelismConfig(\n    sp_backend=\"deepspeed\",\n    sp_size=2,\n    dp_shard_size=2,  # 2D parallelism: SP × DP = 4 GPUs\n    sp_handler=DeepSpeedSequenceParallelConfig(\n        sp_seq_length_is_variable=True,\n        sp_attn_implementation=\"flash_attention_2\",\n    ),\n)\n\ntraining_args = SFTConfig(\n    ...,\n    parallelism_config=parallelism_config,\n    max_length=32768,\n    pad_to_multiple_of=2,  # Must equal sp_size\n    per_device_train_batch_size=1,\n)\n\ntrainer = SFTTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=dataset,\n)\ntrainer.train()\n```\n\n### Key SFTConfig Parameters for Ulysses\n\n| Parameter            | Description                                                                                  |\n| -------------------- | -------------------------------------------------------------------------------------------- |\n| `pad_to_multiple_of` | Must equal `sp_size` to ensure sequence divisibility                                         |\n| `max_length`         | Global sequence length (before splitting across GPUs)                                        |\n| `packing`            | Works well with SP — packing reduces padding waste, especially for variable-length sequences |\n\n### Accelerate Config File\n\nCreate `alst_ulysses_4gpu.yaml`:\n\n```yaml\ncompute_environment: LOCAL_MACHINE\ndistributed_type: DEEPSPEED\nmixed_precision: bf16\nnum_processes: 4\ndeepspeed_config:\n  zero_stage: 3\n  seq_parallel_communication_data_type: bf16\nparallelism_config:\n  parallelism_config_sp_size: 2\n  parallelism_config_sp_backend: deepspeed\n  parallelism_config_dp_shard_size: 2\n  parallelism_config_sp_seq_length_is_variable: true\n  parallelism_config_sp_attn_implementation: flash_attention_2\n```\n\n### Complete Training Command\n\n```bash\naccelerate launch --config_file alst_ulysses_4gpu.yaml \\\n    trl/scripts/sft.py \\\n    --model_name_or_path meta-llama/Llama-3.1-8B \\\n    --dataset_name trl-lib/Capybara \\\n    --max_length 32768 \\\n    --packing \\\n    --pad_to_multiple_of 2 \\\n    --per_device_train_batch_size 1\n```\n\n### Shift Labels Handling\n\nThe SFTTrainer automatically handles pre-shifted labels when Ulysses is enabled:\n\n```python\n# When using SP, labels are pre-shifted by the dataloader adapter\n# The trainer detects this and uses shift_labels directly\nlabels = inputs[\"labels\"] if \"shift_labels\" not in inputs else None\n\n# Loss computation uses the pre-shifted labels\nif \"shift_labels\" in inputs:\n    shift_logits = outputs.logits.contiguous()\n    shift_labels = inputs[\"shift_labels\"]\nelse:\n    shift_logits = outputs.logits[..., :-1, :].contiguous()\n    shift_labels = labels[..., 1:].contiguous()\n```\n\n## Comparing Ulysses and Ring Attention\n\nBoth Ulysses and Ring Attention enable long-context training, but they have different characteristics:\n\n| Aspect                    | Ulysses (DeepSpeed)               | Ring Attention (FSDP2) |\n| ------------------------- | --------------------------------- | ---------------------- |\n| **Parallelism Method**    | Attention head partitioning       | Ring-based KV exchange |\n| **Backend**               | DeepSpeed ZeRO                    | PyTorch FSDP2          |\n| **Attention Support**     | FlashAttention 2/3, SDPA          | SDPA only              |\n| **Communication**         | Two `all-to-all`s per layer       | P2P ring communication |\n| **Comm volume per GPU**   | O(total\\_seq x hidden / sp\\_size) | O(total\\_seq x hidden) |\n| **Sequence Divisibility** | `sp_size`                         | `cp_size * 2`          |\n| **Num Head Constraint**   | `num_heads >= sp_size`            | None                   |\n\n### When to Choose Ulysses vs Ring Attention\n\nSince switching between the two only requires changing the accelerate config, we recommend trying both and comparing performance and memory usage on your specific setup. The main constraint is that Ulysses requires `num_heads >= sp_size`, while Ring Attention has no such limitation.\n\n## Best Practices\n\n### 1. Sequence Length Divisibility\n\nAlways ensure your sequence length is divisible by `sp_size`:\n\n```python\ntraining_args = SFTConfig(\n    pad_to_multiple_of=4,  # For sp_size=4\n    max_length=32768,  # Must be divisible by 4\n)\n```\n\n### 2. Use Flash Attention\n\nFlash Attention 2 provides cleaner output and better performance than SDPA:\n\n```python\nparallelism_config = ParallelismConfig(\n    sp_handler=DeepSpeedSequenceParallelConfig(\n        sp_attn_implementation=\"flash_attention_2\",\n    ),\n)\n```\n\nUse Flash Attention 3 for Hopper and look out for Flash Attention 4 release for Blackwell (FA2 on Blackwell is quite slow).\n\n### 3. Combine with DeepSpeed ZeRO\n\nFor very large models, combine Ulysses with ZeRO Stage 3:\n\n```yaml\ndeepspeed_config:\n  zero_stage: 3\n  offload_optimizer:\n    device: cpu\n```\n\nIf the model is huge, you can offload the params as well by adding to the above:\n\n```yaml\n  offload_param:\n    device: cpu\n```\n\n### 5. Use memory fragmentation-friendly PyTorch allocator\n\nThis environment variable will allow for a longer sequence length:\n\n```bash\nexport PYTORCH_ALLOC_CONF=expandable_segments:True\n```\n\n### 6. 2D Parallelism Configuration\n\nBalance SP and DP for your GPU count:\n\n| GPUs | `sp_size` | `dp_shard_size` | Use Case                                        |\n| ---- | --------- | --------------- | ----------------------------------------------- |\n| 4    | 2         | 2               | Balanced throughput and sequence length         |\n| 4    | 4         | 1               | Maximum sequence length                         |\n| 8    | 2         | 4               | Higher throughput with moderate sequence length |\n| 8    | 4         | 2               | Longer sequences with moderate throughput       |\n\nRemember: `dp_replicate_size × dp_shard_size × sp_size = num_processes`\n\n### 7. Liger-Kernel\n\nIf your desired model architecture is supported by [Liger-Kernel](https://github.com/linkedin/Liger-Kernel), it is fully compatible with Ulysses SP and can be enabled with a single flag:\n\n```python\ntraining_args = SFTConfig(\n    use_liger_kernel=True,\n)\n```\n\nThe main memory saving comes from `FusedLinearCrossEntropy` which avoids materializing the full logits tensor during loss calculation. The savings grow with longer sequences where the logits tensor is larger.\n\nAdditionally, you can enable [`TiledMLP`](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/#tiled-mlp-computation) to further extend sequence length — like `FusedLinearCrossEntropy`, it saves working memory by tiling large matrix operations.\n\n### 8. Token Distribution Across Ranks\n\nYou don't need to 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 reasonably sized dataset, the distribution evens out statistically over training.\n\n## Benchmarks\n\nTo quantify the benefits of Ulysses SP, we trained [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) on the [Gutenberg English](https://huggingface.co/datasets/sedthh/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 attention backend.\n\n### Setup\n\n| Config   | GPUs | SP | DP | Seq Length | Grad Acc | Global Batch |\n| -------- | ---- | -- | -- | ---------- | -------- | ------------ |\n| Baseline | 1    | 1  | 1  | 8K         | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 8K         | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 32K        | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 64K        | 8        | 8            |\n| SP=4     | 4    | 4  | 1  | 96K        | 8        | 8            |\n\nThe benchmark runs in the table above use the same global batch size (8 micro-batches), cosine learning-rate schedule, and seed, so those benchmark loss curves are directly comparable.\n\n### Loss Curve Matching Diagnostics (4 GPU)\n\nTo verify SP-vs-DP loss equivalence, we ran controlled 4-GPU A/B experiments with identical seed, model, optimizer, learning-rate schedule, and data order.\n\n#### Methodology for Fair DP vs SP Comparison\n\nCompared setups:\n\n- `DP=4, SP=1, GAS=1` (baseline)\n- `DP=1, SP=4, GAS=4` (Ulysses SP)\n\nFor fair comparison, `GAS` must scale with `SP`:\n\n- Ulysses SP splits the sequence across `SP` ranks, so each SP rank sees roughly `1/SP` of the sequence tokens per micro-step.\n- If `GAS` is unchanged, each optimizer step in SP aggregates fewer total tokens than the DP baseline.\n- Setting `GAS=SP` keeps effective tokens per optimizer step matched:\n  - DP tokens/step: `dp_world_size * micro_batch * seq_len * GAS = 4 * B * L * 1`\n  - SP tokens/step: `dp_world_size * micro_batch * (L/SP) * GAS * SP_ranks = 1 * B * (L/4) * 4 * 4 = 4 * B * L`\n\n![Canonical loss on Gutenberg for DP=4 vs SP=4](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/gutenberg_canonical_loss_dp4_vs_sp4.png)\\\nOn Gutenberg text (20 steps), canonical loss matches within logging precision between `DP=4,SP=1,GAS=1` and `DP=1,SP=4,GAS=4`.\n\nMeasured over 20 steps on 4 GPUs in controlled equivalence harnesses:\n\n| Harness      | Metric        | DP vs SP setting         | Mean abs diff | Max abs diff |\n| ------------ | ------------- | ------------------------ | ------------- | ------------ |\n| `Trainer`    | `loss`        | DP=4, SP=1 vs DP=1, SP=4 | 0.0054        | 0.0131       |\n| `SFTTrainer` | logged `loss` | DP=4, SP=1 vs DP=1, SP=4 | 0.0811        | 0.0812       |\n| `SFTTrainer` | canonical NLL | DP=4, SP=1 vs DP=1, SP=4 | 0.000004      | 0.000005     |\n\nTakeaway: under matched token budget, SP and non-SP match on canonical token-normalized loss. The remaining difference is in trainer-reported logging (`loss`), not in the underlying cross-entropy objective.\n\n### Memory Reduction\n\n![Peak GPU Memory per Rank](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/memory_usage.png)\\\nSP=4 reduces per-GPU memory by 3.3x at the same sequence length, enabling training at up to 96K tokens on 4× H100 80GB. At 128K, the model OOMs.\n\n| Config       | Seq Length | Peak Memory | Notes                             |\n| ------------ | ---------- | ----------- | --------------------------------- |\n| DP=4 (4 GPU) | 8K         | 22.4 GB     | Baseline — no SP                  |\n| SP=4 (4 GPU) | 8K         | 22.8 GB     | Similar memory at same seq length |\n| SP=4 (4 GPU) | 32K        | 35.0 GB     | 4x longer than DP baseline        |\n| SP=4 (4 GPU) | 64K        | 50.5 GB     | 8x longer than DP baseline        |\n| SP=4 (4 GPU) | 96K        | 66.0 GB     | **12x longer than DP baseline**   |\n| SP=4 (4 GPU) | 128K       | OOM         | Exceeds 80 GB limit               |\n\nAt 8K tokens, DP=4 and SP=4 use nearly the same memory per GPU (~22 GB with ZeRO-3). The advantage of SP is that it enables scaling to much longer sequences: at 96K tokens (12x longer), peak memory is 66 GB — still within the H100's 80 GB capacity. At 128K, the model OOMs, establishing the practical limit for this configuration. DP=4 without SP cannot scale beyond 8K for this model.\n\n### Throughput\n\n![Training Throughput](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ulysses/throughput.png)\\\nLonger sequences with SP process dramatically more tokens per second. SP=4 at 64K achieves 3.7x the throughput of the baseline.\n\n| Config           | Seq Length | Tokens/s | vs Baseline |\n| ---------------- | ---------- | -------- | ----------- |\n| Baseline (1 GPU) | 8K         | 3,633    | —           |\n| SP=4 (4 GPU)     | 8K         | 3,933    | \\~1x        |\n| SP=4 (4 GPU)     | 32K        | 7,733    | **2.1x**    |\n| SP=4 (4 GPU)     | 64K        | 13,396   | **3.7x**    |\n\nAt the same 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 real benefit 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.\n\n> These results use only 4 GPUs with SP=4. With 8 GPUs (SP=8), you can push to even longer sequences — up to 256K+ tokens — or use 2D parallelism (SP=4, DP=2) to combine long-context training with data-parallel throughput.\n\n## Requirements\n\n- HF Accelerate: `deepspeed>=0.18.1 accelerate>=1.12`\n- HF Trainer: `deepspeed>=0.18.1 accelerate>=1.12 transformers>=5.0`\n- HF TRL: `deepspeed>=0.18.1 accelerate>=1.12 transformers>=5.0 trl>=0.18.0`\n\nUse `flash_attention_2` for Ampere GPUs, or `flash_attention_3` for Hopper GPUs. Wait for `flash_attention_4` on Blackwell 🕰.\n\n## Resources\n\n### Documentation\n\n- [Accelerate: Context Parallelism Guide](https://huggingface.co/docs/accelerate/concept_guides/context_parallelism)\n- [TRL: Distributing Training](https://huggingface.co/docs/trl/distributing_training)\n- [DeepSpeed Sequence Parallelism](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism)\n\n### Examples\n\n- [Accelerate ALST Example](https://github.com/huggingface/accelerate/tree/main/examples/alst_ulysses_sequence_parallelism)\n- [TRL Accelerate Configs](https://github.com/huggingface/trl/tree/main/examples/accelerate_configs)\n\n### Papers\n\n- [Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences](https://huggingface.co/papers/2506.13996)\n- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://huggingface.co/papers/2309.14509)\n\n### Related Blog Posts\n\n- [Accelerate ND-Parallel: A Guide to Efficient Multi-GPU Training](https://huggingface.co/blog/accelerate-nd-parallel)\n- [Understanding Ulysses and Ring Attention](https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention)\n- [Enabling Long-Context Training with Sequence Parallelism in Axolotl](https://huggingface.co/blog/axolotl-ai-co/long-context-with-sequence-parallelism-in-axolotl)\n",
  "dateCreated": "2026-10-04T12:50:30.804479567+00:00",
  "date_created": "2026-03-09T00:00:00+00:00",
  "date_modified": "2026-03-09T00:00:00+00:00",
  "date_published": "2026-03-09T00:00:00+00:00",
  "description": "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…",
  "external_url": "https://huggingface.co/blog/ulysses-sp",
  "git": {
    "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e",
    "path": "items/huggingface-co/2026/03/2026-03-09-ulysses-sequence-parallelism-training-with-million-token.md",
    "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/huggingface-co/2026/03/2026-03-09-ulysses-sequence-parallelism-training-with-million-token.md"
  },
  "id": "urn:aggr:item:ae61e2bf024e6021f20842d262a06bc775d0bdb4",
  "isBasedOn": {
    "@id": "https://huggingface.co/blog/ulysses-sp"
  },
  "mainEntity": {
    "@id": "https://huggingface.co/blog/ulysses-sp",
    "@type": "CreativeWork",
    "archivedAt": {
      "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-03-09-ulysses-sequence-parallelism-training-with-million-token/#webpage"
    },
    "dateModified": "2026-03-09T00:00:00+00:00",
    "datePublished": "2026-03-09T00:00:00+00:00",
    "headline": "Ulysses Sequence Parallelism: Training with Million-Token Contexts",
    "timeRequired": "PT13M",
    "url": "https://huggingface.co/blog/ulysses-sp",
    "wordCount": 2814
  },
  "name": "Archived snapshot: Ulysses Sequence Parallelism: Training with Million-Token Contexts",
  "preview": null,
  "reading_minutes": 13,
  "replicated_at": null,
  "source": "huggingface-co",
  "source_name": "Hugging Face - Blog",
  "summary": null,
  "tags": [],
  "temporalCoverage": "2026-10-04T12:50:30.804479567+00:00",
  "title": "Ulysses Sequence Parallelism: Training with Million-Token Contexts",
  "url": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-03-09-ulysses-sequence-parallelism-training-with-million-token/",
  "word_count": 2814
}