Linear Attention Sequence Parallelism

Sequence parallel serves as a prevalent strategy to handle long sequences that exceed the memory limit of a single GPU. However, existing methods do not take advantage of linear attention features, resulting in sub-optimal parallelism efficiency and usability for linear-complexity language models. In this paper, we introduce Linear Attention Sequence Parallel (LASP), an efficient sequence parallel method designed for linear attention-based language models. Specifically, we design an efficient point-to-point communication mechanism to leverage the right-product kernel trick of linear attention, which sharply decreases the communication overhead. We enhance the practical efficiency of LASP by performing kernel fusion and intermediate state caching, making the implementation of LASP hardware-friendly on GPU clusters. Furthermore, we meticulously ensure the compatibility of sequence-level LASP with all types of batch-level data parallel methods, which is vital for distributed training on large clusters with long sequences and large batches. We also discuss the versatility of LASP on other linear-complexity models. Extensive experiments on linear attention-based models are conducted with varying sequence lengths and GPU cluster sizes. LASP scales sequence length up to 4096K using 128x A100 80G GPUs on 1B models, which is 8$\times$ longer than existing methods while being significantly faster.