Sequence Parallelism: Long Sequence Training from System Perspective (2023)

Abstract: Transformer achieves promising results on various 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 algorithm perspective. In this work, we propose sequence parallelism, a memory-efficient parallelism to solve this issue from system perspective instead. Our approach is compatible with most existing parallelisms (e.g., data, pipeline, and tensor parallelism), which means our sequence parallelism makes 4D parallelism possible. More importantly, we no longer require a single device to hold the whole sequence. Besides, using efficient attention with linear complexity, 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). Experiments show that sequence parallelism performs well when scaling with batch size and sequence length. Compared with tensor parallelism, our approach achieved 13.7\times and 3.0\times maximum batch size and sequence length respectively when scaling up to 64 NVIDIA P100 GPUs. With efficient attention, sequence can handle sequence with over 114K tokens, which is over 27\times longer than existing efficient attention works holding the whole sequence on a single device.

## introduction transformer-based language models @xcite @xcite have achieved impressive performance on various natural language understanding and generation tasks (e.g., q&a @xcite , relation extraction (xue et al., 2020b,a; zhou et al., * equal contribution 2020) and dialogue system @xcite ). recently, transformer also achieved promising results on computer vision tasks @xcite @xcite and even on bioinformatics tasks @xcite . these transformer-based models learn powerful context-aware representation by applying self-attention to all pairs of tokens from the input sequence. this mechanism captures longterm dependencies at the token level for sequence modeling. however, self-attention suffers from quadratic memory requirements with respect to sequence length. existing works focusing on long sequence modeling devote to solve this problem from algorithm perspective. that is, these works mainly try to reduce the time and space complexity of attention. in this paper, we focus on solving the long sequence training problem from system perspective. 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 for medical image classification, each image is much larger than it is in usual (e.g., 512×512×512 vs 256×256×3). then, each medical image has much more tokens (i.e., over 512×). each input sequence is much longer than usual. in this case, it is challenging to hold the whole sequence within single gpu. in this paper, we designed and implemented sequence parallelism, which aims at breaking the limitation that we must store the whole sequence in one gpu. the proposed system can train transformerbased models with longer sequences and a larger batch size. specifically, we first split the input sequence into multiple chunks along the sequence dimension and feed each sub-sequence chunk to one corresponding gpu. each gpu thus only holds a part of the full sequence, i.e., a sub-sequence. to apply self-attention to the tokens from different chunks, the main challenge is to compute atten-tion scores and outputs across gpus efficiently. to tackle this problem, we proposed ring self-attention (rsa), which circulates key and value embeddings across gpus in a ring manner. in this case, each device is just required to keep the attention embeddings corresponding to its own subsequence. as a result, our sequence parallelism is memory-efficient, especially for long input sequences. to model long sequences, existing works mainly focus on efficient attention (e.g., @xcite ) with linear instead of quadratic space complexity. in this paper, we aim to solve the long sequence modeling problem from the distributed system perspective. we evaluated our system on both vanilla attention to verify our system is a general solution, and evaluated on efficient attention setting to show the upper bound sequence length. existing pipeline parallelism @xcite and tensor parallelism @xcite ) are designed to cope with a larger model size instead of longer sequences. however, when the sequence is long, the challenge is, existing parallelism must keep the whole sequence on one single device. even if splitting model along hidden and attention-head dimension (i.e., tensor parallelism) or depth dimension (i.e., pipeline parallelism) can still process longer sequences to some extent, the attention-head and depth are much smaller than sequence length (e.g., 12 vs 512), which limits the training scalability and the maximum length of the input sequence. in contrast, our approach splits the whole sequence into multiple devices, enabling it to fit longer input data. in summary, our main contributions are three folds: pipeline parallelism huge deep neural networks @xcite ) have shown their effectiveness on various tasks. however, it is challenging to hold the whole model on one single device due to memory limitations. to overcome this, @xcite proposed pipeline parallelism, model parallelism splitting the model layers into different partitions on separate accelerators. as shown in figure 1a , they split the data along the batch dimension into micro-batches, and each device can process one micro-batch received from the previous device at a time. when the computation is pipelined across micro-batches, pipelining schemes need to ensure that inputs use consistent weight versions for both forward and backward computation to ensure correct weight update and model convergence @xcite . tensor parallelism different from pipeline parallelism which splits models by layer, tensor parallelism (i.e., megatron) @xcite ) introduces tensor splitting, where individual layers of the model are partitioned over multiple devices. similar to our sequence parallelism, tensor parallelism is also designed for transformerbased models. each transformer layer includes a self-attention block and a two-layer multi-layer perceptron (mlp) block. the mlp block can be formalized as: at the second gemm, z 1 and z 2 need to undergo an all-reduce operation to give the final output before the dropout layer in the transformer layer. similarly, megatron splits the tensors in the selfattention layer as well. for multi-head attention, attention heads are split by column and allocated equally to the devices. the linear layer after the self-attention computation is split by row. an allreduce operation is needed at the linear layer output to aggregate attention output from all devices. please refer to megatron @xcite for more details about tensor parallelism. ## sequence parallelism we propose sequence parallelism for training transformer with longer sequences. the overview of sequence parallelism is shown in figure 1c . input sequences are split into multiple chunks and the sub-sequences are fed to different corresponding devices. all devices are holding the same trainable parameters but different sub-sequence input chunks. we will introduce and analyze sequence parallelism in detail below. we use the following notation in this section: ( 1 ## conclusion in this paper, we proposed sequence parallelism for training transformer with longer sequence. sequence parallelism is designed to break the limitation of sequence length on a single device. we have shown that sequence parallelism can handle longer sequence and is more memory-efficient than sota. in particular, sequence parallelism achieves 3.0× maximum sequence length and 13.7× maximum batch size than tensor parallelism when scaling up to 64 gpus. unlike both tensor and pipeline parallelism, sequence parallelism is not limited by the smaller hyper-parameters (e.g., number of attention heads, number of layers). therefore, our sequence parallelism can be adapted as long as the sequence length is divisible by sequence parallel size. with efficient attention, sequence parallelism can handle sequence with over 114k tokens, which is over 27× longer than existing efficient attention works holding the whole sequence on a single device. we used a language model (i.e., bert) to evaluate our system, but it can also be adapted to vision tasks. this work paves the way to process large images @xcite by vit @xcite as a larger image means more patches or longer sequences.