Source-linked AI summary
DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, Yuxiong He
TL;DR
Existing sequence-parallelism systems are limited by memory and communication inefficiency for long-sequence Transformer training. DeepSpeed-Ulysses partitions sequences across GPUs and uses all-to-all communication for attention, achieving up to 2.5x throughput improvements while supporting 4x longer sequences than existing systems.
Problem
Existing sequence-parallelism methods are constrained by memory-communication inefficiency, limiting their scalability to long-sequence large models.
Method
DeepSpeed-Ulysses partitions samples along the sequence dimension and uses all-to-all communication to distribute attention computation across GPUs.
Results
Up to 2.5x throughput improvements and over 10x communication reduction are reported, while supporting 4x longer sequences and sequences exceeding one million tokens.
Takeaways & Limitations
DeepSpeed-Ulysses provides portable sequence parallelism that supports dense and sparse attention and enables model size and sequence length to scale beyond single-GPU memory limits.
Abstract
from arXiv · showhide
Computation in a typical Transformer-based large language model (LLM) can be characterized by batch size, hidden dimension, number of layers, and sequence length. Until now, system works for accelerating LLM training have focused on the first three dimensions: data parallelism for batch size, tensor parallelism for hidden size and pipeline parallelism for model depth or layers. These widely studied forms of parallelism are not targeted or optimized for long sequence Transformer models. Given practical application needs for long sequence LLM, renewed attentions are being drawn to sequence parallelism. However, existing works in sequence parallelism are constrained by memory-communication inefficiency, limiting their scalability to long sequence large models. In this work, we introduce DeepSpeed-Ulysses, a novel, portable and effective methodology for enabling highly efficient and scalable LLM training with extremely long sequence length. DeepSpeed-Ulysses at its core partitions input data along the sequence dimension and employs an efficient all-to-all collective communication for attention computation. Theoretical communication analysis shows that whereas other methods incur communication overhead as sequence length increases, DeepSpeed-Ulysses maintains constant communication volume when sequence length and compute devices are increased proportionally. Furthermore, experimental evaluations show that DeepSpeed-Ulysses trains 2.5x faster with 4x longer sequence length than the existing method SOTA baseline.
1 Introduction
Long-sequence training is increasingly needed for generative AI and scientific applications, but existing parallelism methods do not efficiently scale along the sequence dimension. DeepSpeed-Ulysses addresses this gap with sequence partitioning and all-to-all communication, supporting longer sequences with improved efficiency and portability.
- Long sequences support applications including conversational AI, document summarization, video generation, multimodal reasoning, and scientific discovery.
- Existing data, tensor, and pipeline parallelism methods do not scale along the sequence dimension, while sequence-parallel approaches face memory-communication inefficiency.
- DeepSpeed-Ulysses partitions samples across the sequence dimension and uses all-to-all communication so GPUs compute attention for non-overlapping head subsets.
- 4x longer sequence lengths and sequences exceeding one million tokens are supported compared with existing systems.
- Over 10x communication reduction and up to 2.5x throughput improvement are reported, with sustained throughput above 175 TFlops/GPU.
- The approach supports dense and sparse attention, works with FlashAttention v2, integrates with ZeRO-3, and requires minimal code changes.
2 Background and Related Work
Transformer training uses data, model, and sequence parallelism to distribute computation, but extremely long sequences create intermediate-activation memory and communication challenges. DeepSpeed-Ulysses differs from related sequence-parallel methods through communication efficiency, ZeRO compatibility, attention flexibility, and portability.
- 2.1.1 Transformer Architecture: A Transformer projects input sequences into query, key, and value embeddings before attention and multilayer-perceptron processing.
- 2.1.2 Mode of Parallelism: Data parallelism partitions samples while replicating parameters, but its effectiveness depends on sufficiently large batches and large batches can affect model quality.
- 2.1.2 Mode of Parallelism: Sequence parallelism is orthogonal to data parallelism and ZeRO, and can accelerate fixed-length training or enable progressively longer contexts.
- 2.2 Related Work: Existing parallel methods remain limited by intermediate activation memory for extremely long sequences, while sequence-parallel methods also differ in communication efficiency.
- 2.2 Related Work: ColAI-SP uses ring communication for keys and values, whereas Megatron-LM combines sequence parallelism with tensor parallelism and collective aggregation.
- 2.2 Related Work: DeepSpeed-Ulysses is presented as more communication-efficient, ZeRO-compatible, attention-flexible, and easier to use than the compared sequence-parallel methods.
3 DeepSpeed-Ulysses Core Design
DeepSpeed-Ulysses partitions sequences across devices, uses all-to-all communication around attention, and returns sequence-parallel outputs for subsequent Transformer operators. Its communication volume scales better with sequence parallelism than existing approaches, while supporting general attention implementations and integration with ZeRO-3.
- System Design: DeepSpeed-Ulysses partitions each input sequence across P devices, gathers global QKV with all-to-all communication, and computes attention on the gathered embeddings.After attention, a second all-to-all restores N/P sequence partitions for later Transformer operators.
- System Design: The second all-to-all transforms attention outputs back to sequence-parallel N/P partitions for MLP, matrix multiplication, and layer-normalization modules.
- Communication Analysis: For aggregate all-to-all message size M over P GPUs, the per-link volume is M/P; DS-Sequence therefore communicates 3Nh and Nh for its two collectives.The analysis uses hidden size h, sequence length N, and parallelism degree P.
- Communication Analysis: Megatron-LM incurs 4Nh communication per link, P times larger than DeepSpeed sequence parallelism, and its communication complexity is O(N).DeepSpeed-Ulysses maintains constant communication volume when N and P increase proportionally, whereas Megatron-LM grows linearly with N.
- Memory Efficiency: Because sequence parallelism reduces activation memory but not model-state memory, DeepSpeed-Ulysses integrates with ZeRO-3 across data- and sequence-parallel ranks.ZeRO-3 partitions model states across the combined parallel groups and gathers needed partitions.
- System Design: DeepSpeed-Ulysses supports self-, cross-, and causal attention in dense or sparse forms, including optimized kernels such as FlashAttention.Its attention-centric modular design makes the implementation general across attention variants and kernels.
4 Evaluation
The evaluation measures DeepSpeed-Ulysses across sequence-length scaling, dense and sparse attention throughput, parallel scaling, and convergence. Across these studies, it supports longer sequences and generally outperforms or scales efficiently relative to the compared systems.
- Sequence-length scalability: Sequence length scales linearly with GPU count while maintaining similar computation throughput across appropriate GPU configurations.The strong-scaling study reaches sequence lengths up to 1 million tokens on a 1.2 billion parameter GPT model.
- Dense attention: DeepSpeed-Ulysses consistently outperforms Megatron-LM at shared runnable sequence lengths and can run longer sequences.The comparison covers 7B and 30B dense-attention GPT models on 32 and 64 A100 GPUs, respectively.
- Sparse attention: More than 2x throughput is observed versus Megatron-LM for sparse attention, while ZeRO-3 enables up to 4x longer sequences.The current DeepSpeed throughput is bottlenecked by the local sparse-attention implementation and decreases as sequence length increases.
- Parallel scaling: Parallel-scaling studies show good scaling and high percentages of theoretical peak GPU performance as communication workload increases.The studies vary GPU count at fixed 131,072-token sequences and increase GPU count proportionally with sequence length.
- Convergence: Convergence experiments validate that DeepSpeed-Ulysses does not negatively affect trained-model quality as a system optimization.The study uses a 1.3 billion parameter GPT model at 32K sequence length on 8 A100 GPUs.
5 Conclusion
The conclusion presents DeepSpeed Sequence as a memory- and communication-efficient technology for long-sequence Transformer training. It parallelizes sequences across GPUs and supports dense and sparse Flash attention while allowing model size and sequence length to scale beyond single-GPU memory limits.
- 5 Conclusion: DeepSpeed Sequence enables memory- and communication-efficient long-sequence training across GPUs and supports dense and sparse Flash attention.The conclusion describes sequence parallelism across Transformer components and extension to other AI accelerators.