LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its Hybrid

• URL: http://arxiv.org/abs/2502.07563v1
• Date: Tue, 11 Feb 2025 14:01:39 GMT
• Title: LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its Hybrid
• Authors: Weigao Sun, Disen Lan, Yiran Zhong, Xiaoye Qu, Yu Cheng,
• Abstract summary: 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.
• Score: 25.71221522518279
• License:
• Abstract: Linear sequence modeling approaches, such as linear attention, provide advantages like linear-time training and constant-memory inference over sequence lengths. However, 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, which results in lower computation parallelism, limits their scalability for longer sequences in distributed systems. In this paper, we introduce LASP-2, a new SP method to enhance both communication and computation parallelism when training linear attention transformer models with very-long input sequences. Compared to previous work LASP, LASP-2 rethinks the minimal communication requirement for SP on linear attention layers, reorganizes the whole communication-computation workflow of LASP. In this way, only one single AllGather collective communication is needed on intermediate memory states, whose sizes are independent of the sequence length, leading to significant improvements of both communication and computation parallelism, as well as their overlap. Additionally, we extend LASP-2 to LASP-2H by applying similar communication redesign to standard attention modules, offering an efficient SP solution for hybrid models that blend linear and standard attention layers. Our evaluation on a Linear-Llama3 model, a variant of Llama3 with linear attention replacing standard attention, demonstrates the effectiveness of LASP-2 and LASP-2H. Specifically, LASP-2 achieves training speed improvements of 15.2% over LASP and 36.6% over Ring Attention, with a sequence length of 2048K across 64 GPUs. The Code is released as a part of: https://github.com/OpenSparseLLMs/Linear-MoE.

Related papers

• GP-FL: Model-Based Hessian Estimation for Second-Order Over-the-Air Federated Learning [52.295563400314094]
  Second-order methods are widely adopted to improve the convergence rate of learning algorithms. This paper introduces a novel second-order FL framework tailored for wireless channels.
  arXiv  Detail & Related papers  (2024-12-05T04:27:41Z)
• FusionLLM: A Decentralized LLM Training System on Geo-distributed GPUs with Adaptive Compression [55.992528247880685]
  Decentralized training faces significant challenges regarding system design and efficiency. We present FusionLLM, a decentralized training system designed and implemented for training large deep neural networks (DNNs) We show that our system and method can achieve 1.45 - 9.39x speedup compared to baseline methods while ensuring convergence.
  arXiv  Detail & Related papers  (2024-10-16T16:13:19Z)
• 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)
• Linear Attention Sequence Parallelism [33.06590170649837]
  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.
  arXiv  Detail & Related papers  (2024-04-03T17:33:21Z)
• Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models [20.78813311569383]
  We present Lightning Attention, the first linear attention implementation that enables linear attention to realize its theoretical computational benefits. Specifically, we utilize the conventional attention mechanism for the intra-blocks and apply linear attention kernel tricks for the inter-blocks. Various experiments are conducted on different model sizes and sequence lengths.
  arXiv  Detail & Related papers  (2024-01-09T16:27:28Z)
• 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)
• PlueckerNet: Learn to Register 3D Line Reconstructions [57.20244406275875]
  This paper proposes a neural network based method to solve the problem of Aligning two partially-overlapped 3D line reconstructions in Euclidean space. Experiments on both indoor and outdoor datasets show that the registration (rotation and translation) precision of our method outperforms baselines significantly.
  arXiv  Detail & Related papers  (2020-12-02T11:31:56Z)

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.