arXiv:2502.07563v1 宣布类型: 交叉 摘要: 线性序列建模方法，例如线性注意力，提供了诸如线性时间训练和常量内存推理等优势。然而，现有的序列并行（SP）方法要么未针对线性注意力的正确积计算特性进行优化，要么使用环形通信策略，这导致较低的计算并行性，限制了这些方法在分布式系统中处理较长序列的可扩展性。在本文中，我们引入了LASP-2，这是一种新的SP方法，在使用非常长输入序列训练线性注意力转换器模型时，可以增强通信和计算并行性。与之前的LASP工作相比，LASP-2重新考虑了线性注意力层上SP所需的最小通信需求，重新组织了LASP的整个通信-计算工作流程。这样一来，仅需在中间内存状态上进行一次AllGather集体通信，其规模与序列长度无关，从而在通信和计算并行性以及它们的重叠方面取得了显著改进。此外，我们通过将类似的通信重设计应用于标准注意力模块，将LASP-2扩展为LASP-2H，为混合模型提供了一种高效的SP解决方案，这些混合模型结合了线性和标准注意力层。对Linear-Llama3模型的评估，这是一种使用线性注意力替代标准注意力的Llama3变体，证明了LASP-2和LASP-2H的有效性。具体而言，在64个GPU上使用序列长度2048K时，LASP-2相对于LASP实现了15.2%的训练速度提升，相对于环注意力实现了36.6%的提升。代码作为一部分发布于：https://github.com/OpenSparseLLMs/Linear-MoE。