Elamur Research
☰
home / papers / 2023 / DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models

DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2309.14509

Summary

Standard LLM training parallelism splits work by batch (data parallelism), hidden size (tensor parallelism), or depth (pipeline parallelism). None of these addresses very long input sequences. Earlier sequence-parallel methods, such as ColAI-SP and Megatron-LM sequence parallelism, communicate inefficiently and are hard to integrate. The authors from Microsoft propose DeepSpeed-Ulysses. It shards each sample along the token axis across GPUs. Just before attention, an all-to-all exchange regroups the queries, keys and values so that every GPU holds the whole sequence for its own subset of attention heads. A second all-to-all then restores the sequence sharding. Their analysis shows that per-device communication stays fixed when sequence length and GPU count grow together. With other approaches, it grows with message size. Tests on GPT models with dense and sparse attention on A100 GPUs show the method handles sequences 4x longer than Megatron-LM and runs up to 2.5x faster. The approach adapts to any attention implementation and requires few code changes. This makes million-token contexts practical for long documents, chat histories, and scientific data like genomes.

Key points

  • Communication cost: Ulysses has O(M/P) communication complexity, compared with O(M) for ColAI-SP and Megatron-SP. The authors report more than 10x less communication and over 175 TFLOPs per GPU, which is above 54% of hardware peak. Training on sequences beyond one million tokens is supported.
  • Compatibility: The method works with dense attention, sparse attention, and FlashAttention v2. It combines with ZeRO-3, so very large models and long sequences can be trained together.
  • Main comparison: Experiments used 7B and 30B GPT models on 32 and 64 A100 GPUs. Ulysses beat Megatron-LM on every sequence length both systems could run. With sparse attention it delivered more than 2x the throughput.
  • Scaling tests (GPT-7B): At a fixed 131,072-token sequence, going from 64 to 256 GPUs made each iteration about 3.3x faster. When sequence length (65K→262K) and GPU count were increased together, throughput per GPU stayed near 147–161 TFLOPs.
  • No quality loss: A 1.3B GPT model trained at 32K tokens on 8 A100s showed the same convergence as Megatron-LM, across several ZeRO stages.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: AI & machine learning · Computing & GPUs

Cited by

151 papers of our library cite it (of the 140673 whose bibliographies we read); 5 of them have a page here, 146 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2309.14509 · arXiv 2309.14509

@misc{jacobs2023deepspeed,
  title = {{DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models}},
  author = {Jacobs, Sam Ade and Tanaka, Masahiro and Zhang, Chengming and Zhang, Minjia and Song, Shuaiwen Leon and Rajbhandari, Samyam and He, Yuxiong},
  year = {2023},
  doi = {10.48550/arxiv.2309.14509},
  eprint = {2309.14509},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2309.14509}
}

References

22 references, 2 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. find
  2. Wojciech Kryściński, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: A collection of datasets for long-form narrative summarization, 2022. find
  3. MosaicML. Introducing mpt-7b: A new standard for open-source, commercially usable llms. https://https: //www.mosaicml.com/blog/mpt-7b, 2023. find
  4. Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma.… find our card
  5. Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models, 2023. find
  6. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller,… find
  7. Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. Climax: A foundation model for weather and climate, 2023. find
  8. Maxim Zvyagin, Alexander Brace, Kyle Hippe, Yuntian Deng, Bin Zhang, Cindy Orozco Bohorquez, Austin Clyde, Bharat Kale, Danilo Perez-Rivera, Heng Ma, et al. Genslms: Genome-scale language models reveal sars-cov-2 evolutionary dynamics. bioRxiv, pages 2022–10, 2022. find
  9. Yikuan Li, Ramsey M. Wehbe, Faraz S. Ahmad, Hanyin Wang, and Yuan Luo. Clinical-longformer and clinical-bigbird: Transformers for long clinical sequences. CoRR, abs/2201.11838, 2022a. URL https://arxiv.org/abs/2201. 11838. — [arXiv:2201.11838](https://arxiv.org/abs/2201.11838) ↗
  10. Shang Gao, Mohammed Alawad, M. Todd Young, John Gounley, Noah Schaefferkoetter, Hong Jun Yoon, Xiao-Cheng Wu, Eric B. Durbin, Jennifer Doherty, Antoinette Stroup, Linda Coyle, and Georgia Tourassi. Limitations of transformers on clinical text classification. IEEE Journal of Biomedical and Health Inf… — [doi:10.1109/jbhi.2021.3062322](https://doi.org/10.1109/jbhi.2021.3062322) find
  11. Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. find
  12. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998–6008, 2017. *Figure 8: Convergence evaluation of DeepSpeed-Ulysses with different Z… find
  13. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–16. IEEE, 2020. find
  14. Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’21, 2021. find
  15. Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836, 2016. find
  16. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019. find
  17. Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Granger, Phil Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline parallelism for dnn training. In ACM Symposium on Operating Systems Principles (SOSP 2019), October 2019. find
  18. Yanping Huang, Yonglong Cheng, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. ArXiv, abs/1811.06965, 2018. find
  19. Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline-parallel dnn training. In International Conference on Machine Learning, pages 7937–7947. PMLR, 2021. find
  20. DeepSpeed Team and Rangan Majumder. DeepSpeed: Extreme-scale model training for everyone. https://www. microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/, 2020. find
  21. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:22… find our card
  22. Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, and Yuxiong He. Zero++: Extremely efficient collective communication for giant model training, 2023. Tal Ben-Nun and Torsten Hoefler. Demystifying parallel and distributed deep learning:… find