LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its
  Hybrid
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.
67ad7929dc2968691c241147
Weigao Sun
claimed_verified
2025-02-13T08:21:37.445Z
6246bb33da617c00b48e4d92
/avatars/0304a9f6eb7f5dee4d933d03222f94e9.svg
Weigao Sun
user
weigao266
67ad7929dc2968691c241148
Disen Lan
claimed_verified
2025-02-13T08:21:33.746Z
66ea643899af9ac3463639b1
/avatars/252d470e761a57834dee3dbc60dfefed.svg
Disen Lan
user
landisen
67ad7929dc2968691c241149
Yiran Zhong
claimed_verified
2025-04-04T07:10:13.378Z
64c525e4d68946edad6c7067
/avatars/1b108661634af602717a4ab4b66a151f.svg
Yiran Zhong
user
IanZhong
67ad7929dc2968691c24114a
Xiaoye Qu
67ad7929dc2968691c24114b
Yu Cheng
claimed_verified
2025-06-02T07:53:34.612Z
67017abfe4d49b157ac534d9
/avatars/997e1b9f54b27a7728a9d4abfee4ba91.svg
Yu Cheng
user
ych133
LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its
  Hybrid
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.
LASP-2 is a new sequence parallelism method that improves communication and computation parallelism for training linear attention transformer models with very-long input sequences, demonstrating significant speed improvements over previous methods.
linear attention sequence parallelism AllGather communication parallelism computation parallelism LASP LASP-2 LASP-2H Linear-Llama3 Llama3