Elamur Research⌕
☰
home / papers / 2023 / Sequence Parallelism: Long Sequence Training from System Perspective

Sequence Parallelism: Long Sequence Training from System Perspective

Summary

Self-attention memory grows quadratically with input length. Standard training setups also need each GPU to hold the entire sequence, which caps the usable context. Most prior work attacks this algorithmically with cheaper attention variants. This paper takes a systems approach instead. It introduces sequence parallelism, which cuts each input along the token dimension and assigns one chunk per GPU. All GPUs keep identical model weights. To compute attention across chunks, the authors propose Ring Self-Attention (RSA): key embeddings and then value embeddings are passed around the devices in a ring, so each GPU only stores the embeddings for its own sub-sequence. The method is written in plain PyTorch and works alongside data, pipeline and tensor parallelism, enabling "4D parallelism." On BERT with up to 64 NVIDIA P100 GPUs, it allows much larger batches and longer sequences than Megatron-style tensor parallelism. Combined with linear-complexity attention, it reaches sequences beyond 114K tokens. This matters for domains with inherently long inputs, such as large medical images, ViT on high-resolution data, or AlphaFold-like models.

Key points

  • Compared with tensor parallelism (Megatron) on 64 P100 GPUs, sequence parallelism reaches 13.7× larger maximum batch size and 3.0× longer maximum sequence length.
  • With efficient attention (e.g., Linformer-style projection) on 32 P100 GPUs, it handles sequences over 114K tokens. That is more than 27× longer than sparse/efficient attention works such as BigBird and Linformer, which keep the whole sequence on one device.
  • In weak scaling with pipeline size fixed at 8, raising the global batch from 64 to 512 kept peak memory nearly flat (about 8477–8491 MB). Tensor parallelism grew from 8477 MB to out-of-memory at parallel size 8.
  • The only requirement is that sequence length be divisible by the sequence-parallel size. Unlike tensor or pipeline parallelism, it is not bounded by the number of attention heads or layers.
  • Limitations:
    • Communication overhead is higher than with plain data parallelism.
    • Autoregressive decoding at inference time for unidirectional models cannot be parallelized well across devices.

Interesting ideas

  • I1. Choose the split axis by size. Partition along whichever dimension is largest (sequence length, often hundreds or thousands) rather than heads or layers (often around 12). The same reasoning helps choose a sharding axis for any distributed workload.
  • I2. Ring-style exchange of keys and values. Each device computes partial attention scores against the remote chunks it receives, so the full K/V matrices never sit in one place. This "rotate the data past stationary queries" pattern could apply to other all-pairs computations, such as similarity search over large corpora.
  • I3. Orthogonal, stackable parallelism. Designing the method to compose with data, pipeline and tensor parallelism, and with memory optimizers like ZeRO, makes it a building block rather than a competing framework. This is a good principle for systems tools.
  • I4. Combine systems and algorithmic fixes. Pairing distributed splitting with linear-complexity attention multiplies the gains, giving near-ideal scaling in sequence length. Other memory bottlenecks could likely benefit from the same pairing.
  • I5. Report weak scaling with peak memory alongside throughput. Tracking both shows clearly whether a parallel scheme truly removes per-device memory growth. This is a simple evaluation protocol worth reusing when benchmarking long-context training setups.

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: Natural language processing

Cited by

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

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2023.acl-long.134

@inproceedings{li2023sequence,
  title = {{Sequence Parallelism: Long Sequence Training from System Perspective}},
  author = {Li, Shenggui and Xue, Fuzhao and Baranwal, Chaitanya and Li, Yongbin and You, Yang},
  booktitle = {Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2023},
  doi = {10.18653/v1/2023.acl-long.134},
  url = {https://doi.org/10.18653/v1/2023.acl-long.134}
}

References

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

  1. Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. arXiv preprint arXiv:2005.14165. find
  2. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. find
  3. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. 2020. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929. find
  4. Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rihawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. 2020. Prottrans: Towards cracking the language of life’s code through self-supervised deep learning and high performance computing. arXiv p… find
  5. William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961. find
  6. Le Hou, Youlong Cheng, Noam Shazeer, Niki Parmar, Yeqing Li, Panagiotis Korfiatis, Travis M Drucker, Daniel J Blezek, and Xiaodan Song. 2019. High resolution medical image analysis with spatial partitioning. arXiv preprint arXiv:1909.03108. find
  7. Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. 2018. Gpipe: Efficient training of giant neural networks using pipeline parallelism. arXiv preprint arXiv:1811.06965. find
  8. John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. 2021. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583–589. find our card
  9. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668. find
  10. Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. 2021. Efficient large-scale language model training on gpu clusters. arXiv preprint arXiv:2104.04473. find
  11. Jinjie Ni, Tom Young, Vlad Pandelea, Fuzhao Xue, Vinay Adiga, and Erik Cambria. 2021. Recent advances in deep learning based dialogue systems: A systematic survey. arXiv preprint arXiv:2105.04387. — [arXiv:2105.04387](https://arxiv.org/abs/2105.04387) find
  12. Chen Qu, Liu Yang, Minghui Qiu, W Bruce Croft, Yongfeng Zhang, and Mohit Iyyer. 2019. Bert with history answer embedding for conversational question answering. In Proceedings of the 42nd International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 1133– 1136. find
  13. Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners. find
  14. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67. find
  15. Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. 2021. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. arXiv preprint arXiv:2104.07857. find
  16. Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. 2020. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 3505–3506. find
  17. Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. 2021. Zero-offload: Democratizing billion-scale model training. find our card
  18. Qin Wang, Boyuan Wang, Zhenlei Xu, Jiaxiang Wu, Peilin Zhao, Zhen Li, Sheng Wang, Junzhou Huang, and Shuguang Cui. 2021. Pssm-distil: Protein secondary structure prediction (pssp) on low-quality pssm by knowledge distillation with contrastive learning. find
  19. Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768. find
  20. Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, et al. 2021. Gspmd: General and scalable parallelization for ml computation graphs. arXiv preprint arXiv:2105.04663. — [arXiv:2105.04663](https://arxiv.org/abs/2105.04663) find
  21. Fuzhao Xue, Aixin Sun, Hao Zhang, and Eng Siong Chng. 2020a. An embarrassingly simple model for dialogue relation extraction. arXiv preprint arXiv:2012.13873. — [arXiv:2012.13873](https://arxiv.org/abs/2012.13873) find
  22. Fuzhao Xue, Aixin Sun, Hao Zhang, and Eng Siong Chng. 2020b. Gdpnet: Refining latent multi-view graph for relation extraction. arXiv preprint arXiv:2012.06780. — [arXiv:2012.06780](https://arxiv.org/abs/2012.06780) find
  23. Zekun Yang, Noa Garcia, Chenhui Chu, Mayu Otani, Yuta Nakashima, and Haruo Takemura. 2020. Bert representations for video question answering. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 1556–1565. find
  24. Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33. find our card
  25. Hao Zhang, Aixin Sun, Wei Jing, Liangli Zhen, Joey Tianyi Zhou, and Rick Siow Mong Goh. 2021. Natural language video localization: A revisit in span-based question answering framework. IEEE Transactions on Pattern Analysis and Machine Intelligence. find
  26. Hao Zhang, Aixin Sun, Wei Jing, and Joey Tianyi Zhou. 2020. Span-based localizing network for natural language video localization. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6543–6554, Online. Association for Computational Linguistics. find
  27. Wenxuan Zhou, Kevin Huang, Tengyu Ma, and Jing Huang. 2020. Document-level relation extraction with adaptive thresholding and localized context pooling. arXiv preprint arXiv:2010.11304. — [arXiv:2010.11304](https://arxiv.org/abs/2010.11304) find
  28. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053. find

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/li-2023-sequence-parallelism-long-sequence-training/