Sequence Parallelism: Long Sequence Training from System Perspective

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 2391–2404 July 9-14, 2023 ©2023 Association for Computational Linguistics Sequence Parallelism: Long Sequence Training from System Perspective Shenggui Li, Fuzhao Xue∗, Chaitanya Baranwal, Yongbin Li, Yang You School of Computing, National University of Singapore somerlee.9@gmail.com, f.xue@u.nus.edu Abstract Transformer achieves promising results on vari- ous tasks. However, self-attention suffers from quadratic memory requirements with respect to the sequence length. Existing work focuses on reducing time and space complexity from an al- gorithm perspective. In this work, we propose sequence parallelism, a memory-efficient paral- lelism to solve this issue from system perspec- tive instead. Our approach is compatible with most existing parallelisms (e.g., data, pipeline, and tensor parallelism), which means our se- quence parallelism makes 4D parallelism pos- sible. More importantly, we no longer require a single device to hold the whole sequence. Be- sides, using efficient attention with linear com- plexity, our sequence parallelism enables us to train transformer with infinite long sequence. Specifically, we split the input sequence into multiple chunks and feed each chunk into its corresponding device (i.e., GPU). To compute the attention output, we integrated ring-style communication with self-attention calculation and proposed Ring Self-Attention (RSA). Ex- periments show that sequence parallelism per- forms well when scaling with batch size and sequence length. Compared with tensor paral- lelism, our approach achieved 13.7× and 3.0× maximum batch size and sequence length re- spectively when scaling up to 64 NVIDIA P100 GPUs. With efficient attention, sequence can handle sequence with over 114K tokens, which is over 27× longer than existing efficient atten- tion works holding the whole sequence on a single device. 1 Introduction Transformer-based language models (Radford et al., 2019; Brown et al., 2020; Devlin et al., 2018) have achieved impressive performance on various nat- ural language understanding and generation tasks (e.g., Q&A (Qu et al., 2019; Yang et al., 2020), relation extraction (Xue et al., 2020b,a; Zhou et al., ∗Equal Contribution 2020) and dialogue system (Ni et al., 2021)). Re- cently, Transformer also achieved promising results on computer vision tasks (Dosovitskiy et al., 2020; Zhang et al., 2020, 2021) and even on bioinfor- matics tasks (Elnaggar et al., 2020; Wang et al., 2021). These Transformer-based models learn powerful context-aware representation by apply- ing self-attention to all pairs of tokens from the input sequence. This mechanism captures long- term dependencies at the token level for sequence modeling. However, self-attention suffers from quadratic memory requirements with respect to se- quence length. Existing works focusing on long se- quence modeling devote to solve this problem from algorithm perspective. That is, these works mainly try to reduce the time and space complexity of at- tention. In this paper, we focus on solving the long sequence training problem from system perspec- tive. Existing system requires us to hold the whole sequence in one GPU, which limits the length of input sequence. Unfortunately, the long sequence is common in real-world applications. For instance, when we train Transformer …

Content truncated to 500 words