Linear Attention Sequence Parallelism

• URL: http://arxiv.org/abs/2404.02882v2
• Date: Sun, 12 Jan 2025 12:01:47 GMT
• Title: Linear Attention Sequence Parallelism
• Authors: Weigao Sun, Zhen Qin, Dong Li, Xuyang Shen, Yu Qiao, Yiran Zhong,
• Abstract summary: We introduce Linear Attention Sequence Parallelism (LASP) for linear attention-based transformer models. LASP takes advantage of the right-product kernel trick of linear attention, which sharply decreases the communication overhead. LASP scales sequence length up to 4096K on 128 GPUs, which is 8$times$ longer than existing SP methods.
• Score: 33.06590170649837
• License:
• Abstract: Sequence parallelism (SP) serves as a prevalent strategy to handle long sequences that exceed the memory limit of a single device. However, for linear sequence modeling methods like linear attention, existing SP approaches do not take advantage of their right-product-first feature, resulting in sub-optimal communication efficiency and usability. In this paper, we introduce Linear Attention Sequence Parallelism (LASP), an efficient SP approach designed for linear attention-based transformer models. Specifically, we design an efficient point-to-point ring-style communication mechanism to leverage the right-product kernel trick of linear attention, which sharply decreases the communication overhead, comparing with existing SP methods. We enhance the computation efficiency of LASP by performing kernel fusion and intermediate state caching, making the implementation of LASP hardware-friendly on GPUs. 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 very-long sequences. We also discuss the generalization of LASP on other linear sequence modeling methods. Extensive experiments on linear attention-based models are conducted with varying sequence lengths from 2K to 4096K. LASP scales sequence length up to 4096K on 128 GPUs, which is 8$\times$ longer than existing SP methods. The code is available at https://github.com/OpenNLPLab/LASP.

Related papers

• LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its Hybrid [25.71221522518279]
  Linear sequence modeling approaches, such as linear attention, provide advantages like linear-time training and constant-memory inference over sequence lengths. Existing sequence parallelism (SP) methods are either not optimized for the right-product-first feature of linear attention or use a ring-style communication strategy. We introduce LASP-2, a new SP method to enhance both communication and computation parallelism when training linear attention transformer models.
  arXiv  Detail & Related papers  (2025-02-11T14:01:39Z)
• Parallel Sequence Modeling via Generalized Spatial Propagation Network [80.66202109995726]
  Generalized Spatial Propagation Network (GSPN) is a new attention mechanism for optimized vision tasks that inherently captures 2D spatial structures. GSPN overcomes limitations by directly operating on spatially coherent image data and forming dense pairwise connections through a line-scan approach. GSPN achieves superior spatial fidelity and state-of-the-art performance in vision tasks, including ImageNet classification, class-guided image generation, and text-to-image generation.
  arXiv  Detail & Related papers  (2025-01-21T18:56:19Z)
• FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism [33.23902060961886]
  Existing sequence parallelism methods assume homogeneous sequence lengths (i.e., all input sequences are equal in length) and therefore leverages a single, static scattering strategy for all input sequences. We show that the sequence lengths in LLM training corpora exhibit substantial variability, often following a long-tail distribution. We propose a Heterogeneous-adaptive sequence parallelism method to address this problem.
  arXiv  Detail & Related papers  (2024-12-02T14:16:03Z)
• Mini-Sequence Transformer: Optimizing Intermediate Memory for Long Sequences Training [78.93900796545523]
  Mini-Sequence Transformer (MsT) is a methodology for highly efficient and accurate LLM training with extremely long sequences. MsT partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. integrated with the huggingface library, MsT successfully extends the maximum context length of Qwen, Mistral, and Gemma-2 by 12-24x.
  arXiv  Detail & Related papers  (2024-07-22T01:52:30Z)
• Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences [60.489682735061415]
  We propose CHELA, which replaces state space models with short-long convolutions and implements linear attention in a divide-and-conquer manner. Our experiments on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method.
  arXiv  Detail & Related papers  (2024-06-12T12:12:38Z)
• USP: A Unified Sequence Parallelism Approach for Long Context Generative AI [1.973144426163543]
  Sequence parallelism (SP) is becoming key to unlocking the long-context capabilities of generative AI models. This paper investigates the state-of-the-art SP approaches, i.e. DeepSpeed-Ulysses and Ring-Attention, and proposes a unified SP approach. We achieved 47% MFU on two 8xA800 nodes using SP for the LLAMA3-8B model training using sequence length 208K.
  arXiv  Detail & Related papers  (2024-05-13T13:08:02Z)
• AcceleratedLiNGAM: Learning Causal DAGs at the speed of GPUs [57.12929098407975]
  We show that by efficiently parallelizing existing causal discovery methods, we can scale them to thousands of dimensions. Specifically, we focus on the causal ordering subprocedure in DirectLiNGAM and implement GPU kernels to accelerate it. This allows us to apply DirectLiNGAM to causal inference on large-scale gene expression data with genetic interventions yielding competitive results.
  arXiv  Detail & Related papers  (2024-03-06T15:06:11Z)
• Gated Linear Attention Transformers with Hardware-Efficient Training [60.670102007737476]
  This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability. We then generalize this algorithm to a more expressive variant of linear attention with data-dependent gates. When used as a replacement for the standard attention layer in Transformers, the resulting gated linear attention Transformer is found to perform competitively.
  arXiv  Detail & Related papers  (2023-12-11T18:51:59Z)
• DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models [34.74093040678323]
  We introduce DeepSpeed-Ulysses, a novel, portable and effective methodology for enabling highly efficient and scalable LLM training. DeepSpeed-Ulysses at its core partitions input data along the sequence dimension and employs an efficient all-to-all collective communication for attention. Experiments show that DeepSpeed-Ulysses trains 2.5x faster with 4x longer sequence length than the existing method SOTA baseline.
  arXiv  Detail & Related papers  (2023-09-25T20:15:57Z)
• DBA: Efficient Transformer with Dynamic Bilinear Low-Rank Attention [53.02648818164273]
  We present an efficient yet effective attention mechanism, namely the Dynamic Bilinear Low-Rank Attention (DBA) DBA compresses the sequence length by input-sensitive dynamic projection matrices and achieves linear time and space complexity. Experiments over tasks with diverse sequence length conditions show that DBA achieves state-of-the-art performance.
  arXiv  Detail & Related papers  (2022-11-24T03:06:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.