arXiv:2502.07563cs.LGcs.AI2025-02被引 12

提出新并行方法LASP-2,让长序列线性注意力训练更快更高效。

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

  • 重新设计通信流程,仅需一次AllGather,与序列长度无关。
  • 在64张GPU上处理2048K序列时,训练速度比LASP快15.2%。
  • 适用于混合模型,适合做超长序列生成的开发者参考。

线性序列建模方法(如线性注意力)在训练时间上呈线性、推理内存恒定,优于传统方法。但现有序列并行(SP)策略或未适配线性注意力的右乘优先特性,或采用环形通信,导致计算并行度低,限制了长序列在分布式系统中的扩展性。本文提出LASP-2,一种新型SP方法,提升长序列线性注意力模型训练时的通信与计算并行性。相比先前的LASP,LASP-2重构了通信-计算工作流,仅需一次全归约(AllGather),其数据规模与序列长度无关,显著提升并行效率及重叠度。此外,将该设计扩展至标准注意力模块,形成LASP-2H,支持线性与标准注意力混合模型。在包含线性注意力的Linear-Llama3模型上评估显示:在64张GPU上处理2048K序列时,LASP-2相比LASP提速15.2%,相比环形通信提速36.6%。代码已开源:https://github.com/OpenSparseLLMs/Linear-MoE。

原文摘要 · Abstract (English)

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.

序列并行线性注意力长序列训练分布式训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。