title: LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its
    Hybrid
summary: 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.
paper_authors._id: 67ad7929dc2968691c241147
paper_authors.name: Weigao Sun
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-02-13T08:21:37.445Z
paper_authors._id: 6246bb33da617c00b48e4d92
paper_authors.avatarUrl: /avatars/0304a9f6eb7f5dee4d933d03222f94e9.svg
paper_authors.fullname: Weigao Sun
paper_authors.type: user
paper_authors.user: weigao266
paper_authors._id: 67ad7929dc2968691c241148
paper_authors.name: Disen Lan
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-02-13T08:21:33.746Z
paper_authors._id: 66ea643899af9ac3463639b1
paper_authors.avatarUrl: /avatars/252d470e761a57834dee3dbc60dfefed.svg
paper_authors.fullname: Disen Lan
paper_authors.type: user
paper_authors.user: landisen
paper_authors._id: 67ad7929dc2968691c241149
paper_authors.name: Yiran Zhong
paper_authors._id: 67ad7929dc2968691c24114a
paper_authors.name: Xiaoye Qu
paper_authors._id: 67ad7929dc2968691c24114b
paper_authors.name: Yu Cheng
paper_title: LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its
    Hybrid
paper_summary: 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.
paper_ai_keywords: linear attention sequence parallelism (SP) right-product-first LASP-2 AllGather communication parallelism computation parallelism LASP-2H Linear-Llama3 Llama3