Source-linked AI summary
Sequence Parallelism: Long Sequence Training from System Perspective
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, Yang You
TL;DR
Transformer self-attention has quadratic memory requirements with sequence length, and existing work primarily addresses long-sequence modeling algorithmically. The paper distributes sequence chunks across GPUs and introduces Ring Self-Attention for cross-device attention computation. Compared with tensor parallelism on 64 NVIDIA P100 GPUs, it achieves 13.7× maximum batch size and 3.0× maximum sequence length.
Problem
Self-attention’s quadratic memory requirements with sequence length make long-sequence training difficult, while existing approaches mainly target algorithmic reductions in time and space complexity.
Method
Sequence parallelism splits input sequences across devices, while Ring Self-Attention circulates key and value embeddings among GPUs to compute distributed attention.
Results
13.7× maximum batch size and 3.0× maximum sequence length were achieved versus tensor parallelism when scaling to 64 NVIDIA P100 GPUs.
Takeaways & Limitations
Sequence parallelism removes the requirement that one device hold the whole sequence and is compatible with data, pipeline, and tensor parallelism.
Takeaways & Limitations
The analysis focuses mainly on memory usage and communication cost, and assumes Adam as the optimizer for the memory comparison.
Abstract
from arXiv · showhide
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 method to help us break input sequence length limitation and train with longer sequences on GPUs efficiently. Our approach is compatible with most existing parallelisms (e.g. data parallelism, pipeline parallelism 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. That is, with sparse attention, 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 sparse attention, sequence can handle sequence with over 114K tokens, which is over $27\times$ longer than existing sparse attention works holding the whole sequence on a single device.
1 Introduction
The paper addresses self-attention’s quadratic memory growth by distributing sequence chunks across GPUs rather than requiring one device to store the whole sequence. It introduces Ring Self-Attention and reports larger supported batch sizes and sequence lengths than tensor parallelism.
- Self-attention requires quadratic memory with sequence length, while existing long-sequence work mainly addresses the problem algorithmically.
- Sequence parallelism splits the input sequence into chunks, assigns each chunk to a GPU, and stores only the corresponding subsequence on each device.The devices share trainable parameters while holding different sequence chunks.
- Ring Self-Attention circulates key and value embeddings across GPUs in a ring to compute attention outputs for distributed sequence chunks.
- Sequence parallelism is compatible with data, pipeline, and tensor parallelism, enabling their integration into 4D parallelism.
- 13.7× larger maximum batch size and 3.0× maximum sequence length were achieved versus tensor parallelism when scaling to 64 NVIDIA P100 GPUs.
2 Background
This section reviews self-attention and established pipeline and tensor parallelism approaches. Pipeline parallelism partitions model layers and micro-batches, whereas tensor parallelism partitions layer tensors but still requires communication to aggregate outputs.
- Self-attention: Self-attention forms query, key, and value embeddings for every token and computes each token’s attention against all tokens in the input.The embeddings are combined into matrices Q, K, and V; dk denotes key dimension.
- Pipeline parallelism: Pipeline parallelism partitions model layers across accelerators and processes batch micro-batches through the resulting stages.Consistent weight versions are required across forward and backward computation for correct updates and convergence.
- Tensor parallelism: Tensor parallelism partitions individual model layers across devices, including splitting MLP weight matrices A and B along columns and rows.The MLP uses GeLU between its two matrix multiplications.
- Tensor parallelism: Tensor parallelism requires an all-reduce after the second MLP GEMM to produce the final output before dropout.
- Tensor parallelism: In multi-head attention, tensor parallelism splits attention heads and uses an all-reduce to aggregate attention outputs from devices.
3 Sequence parallelism
Sequence parallelism distributes sequence chunks across devices and uses Ring Self-Attention to compute attention outputs across those chunks. Compared with tensor parallelism, it targets lower memory use for longer sequences and larger batches while maintaining comparable communication overhead.
- Sequence parallelism: Sequence parallelism splits input sequences into chunks, assigning each sub-sequence to a different device while replicating trainable parameters.Each device holds only the attention embeddings for its own sub-sequence.
- Ring Self-Attention: Ring Self-Attention circulates key embeddings N −1 times so each local query computes scores against all sequence chunks.Devices first calculate local partial scores, then repeatedly receive keys from the previous device.
- Ring Self-Attention: RSA then circulates value embeddings to combine each device’s attention scores with all values and produce local attention outputs.The second communication stage mirrors key exchange but transmits values instead.
- Memory modeling: BL > 32H is the condition under which sequence parallelism is more memory-efficient than tensor parallelism in MLP blocks.The comparison considers matrices stored under the two parallelism strategies.
- Memory modeling: BL > 16AZ is the condition under which sequence parallelism is more memory-efficient in the multi-head attention block.RSA uses 2 all-reduce-equivalent operations in forward and 4 in backward, while the extra cost can be offset by no MLP communication.
- Communication cost: Sequence parallelism has the same communication overhead as Megatron tensor parallelism and can avoid one all-gather per pipeline stage.Its forward transfer is 2(N −1) ∗B ∗Z ∗(L/N) ∗A and backward transfer is 6(N −1) ∗B ∗Z ∗(L/N) ∗A.
4 Experiments
Experiments evaluate sequence parallelism on BERT models using P100 GPUs, measuring memory, batch size, throughput, and sparse-attention memory usage.
- Setup: Experiments use BERT Base and BERT Large on Piz Daint, with one 16GB P100 GPU per compute node.The nodes are connected by a high-bandwidth network.
- Sparse attention: Table 3 reports sparse-attention block memory usage, with K denoting Linformer’s projection dimension.The table specifically concerns memory usage in the sparse attention block.
N Ring-AV (B, Z, L
The experiments compare sequence parallelism with tensor parallelism across batch-size, pipeline-size, sequence-length, and weak-scaling settings.
- Maximum batch size: 13.7× larger batch size is achieved by sequence parallelism than Megatron on 64 GPUs versus 12 GPUs for BERT Base.Sequence parallelism also supports a larger parallel size while maintaining comparable throughput at equal parallel size.
- Pipeline scaling: Sequence parallelism achieves higher throughput as pipeline stages increase because sub-sequence activations require no splitting or all-gather between pipeline stages.Megatron holds the full sequence on each device and incurs additional communication for pipeline transfers.
- Maximum sequence length: Around 3× maximum sequence length is achieved on BERT Base when scaling to 64 GPUs, while 16 GPUs still provide 1.4× the tensor-parallel length.Tensor parallelism is limited by splitting across the number of attention heads, whereas sequence parallelism splits the sequence into chunks.
- Sequence-length upper bound: Sparse-attention memory terms containing sequence length L are divided by the number of devices N under sequence parallelism.The evaluation adapts Linformer and uses 32 P100 GPUs with data and pipeline parallel sizes set to 1.
- Weak scaling: Weak scaling keeps memory usage almost constant as global batch size increases, while sequence-length scaling uses less memory with comparable throughput.These results compare sequence parallelism against tensor parallelism with pipeline parallelism fixed at 8.
5 Discussion
The discussion distinguishes sequence parallelism from related distributed-training systems and explains why Megatron is used as the experimental baseline.
- Comparison scope: DeepSpeed is not a direct experimental baseline because it optimizes data-parallel memory footprint, whereas sequence parallelism optimizes a different dimension.The paper states that the two methods are compatible and orthogonal, as with DeepSpeed and Megatron.
- Comparison scope: GShard and GSPMD partition model parameters using TensorFlow’s static computation graph, while sequence parallelism is a PyTorch dynamic-graph tool for longer sequences.The computation-paradigm difference makes GShard and GSPMD unsuitable as baselines in these experiments.
6 Conclusion
The paper concludes that sequence parallelism breaks single-device sequence-length limits and supports longer Transformer training with improved memory scaling.
- Conclusion: 3.0× maximum sequence length and 13.7× maximum batch size are achieved versus tensor parallelism when scaling to 64 GPUs.These are the headline comparisons reported in the conclusion.
- Conclusion: Sequence parallelism requires sequence length to be divisible by the sequence-parallel size, rather than depending on smaller model hyperparameters such as attention heads or layers.This condition defines the stated adaptability boundary.
- Conclusion: Over 114K tokens can be handled with sparse attention, reported as over 27× longer than sparse-attention work keeping the whole sequence on one device.The paper evaluates the system with BERT and notes possible adaptation to vision tasks.
Checklist
The checklist records affirmative responses on contributions and scope, experiment reproducibility details, and use or release of existing assets. It also records that several ethics and participant-related items were not applicable.
- The paper reports that it described its contributions, scope, and limitations.
- The authors marked reproducibility requirements as satisfied, including code or data access, training details, and compute-resource reporting.
- The paper reports citing existing assets, documenting their licenses, and providing new assets.
- The checklist marks participant risks, compensation, and personally identifiable information as not applicable or addressed as specified.
B Convergence performance
The paper evaluates convergence performance by comparing Megatron with sequence parallelism on BERT Large using the Wikipedia development set. The supplied passages describe the evaluation setup but do not state its outcome.
- Convergence was evaluated on Wikipedia’s development set every 1k iterations during 50k-step BERT Large training.Megatron and the proposed model used default Megatron hyperparameters and parallel size 4 without pipeline parallelism.
C Scaling with sequence/tensor parallelism
Sequence parallelism scales favorably against tensor parallelism for BERT Large, supporting substantially larger batches while maintaining comparable throughput at equal parallel size.
- 2.7 times larger batch size was achieved for BERT Large on 16 GPUs.
- 10.2 times larger batch size was achieved on 64 GPUs than with tensor parallelism on 16 GPUs.
- Sequence parallelism achieved comparable throughput with the same parallel size and improved performance at larger parallel sizes.
D Scaling with pipeline parallelism
For BERT Large, sequence parallelism supports higher maximum batch size and better throughput with more pipeline stages, while also extending maximum sequence length through sequence chunking.
- Sequence parallelism achieved a higher maximum batch size than tensor parallelism for BERT Large.
- Sequence parallelism performed better on throughput when using more pipeline stages.
- Around 2× maximum sequence length was achieved on BERT Large when scaling to 64 GPUs.The sequence was split into multiple chunks, enabling better scaling.