LASP-2: Переосмысление последовательной параллельности для линейного внимания и его гибридного подхода

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

February 11, 2025

Авторы: Weigao Sun, Disen Lan, Yiran Zhong, Xiaoye Qu, Yu Cheng

cs.AI

Аннотация

Подходы к моделированию линейной последовательности, такие как линейное внимание, обеспечивают преимущества, такие как обучение за линейное время и вывод с постоянной памятью по длине последовательности. Однако существующие методы параллелизма последовательности (SP) либо не оптимизированы для функции линейного внимания к правому произведению, либо используют стратегию кольцевого обмена информацией, что приводит к снижению параллелизма вычислений и ограничивает их масштабируемость для более длинных последовательностей в распределенных системах. В данной статье мы представляем LASP-2, новый метод SP, который улучшает как коммуникационный, так и вычислительный параллелизм при обучении моделей трансформера с линейным вниманием на очень длинных входных последовательностях. По сравнению с предыдущей работой LASP, LASP-2 переосмысливает минимальные требования к коммуникации для SP на слоях линейного внимания, переорганизует весь рабочий процесс коммуникации и вычислений LASP. Таким образом, требуется всего одно собирающееся общее общение AllGather на промежуточных состояниях памяти, размеры которых не зависят от длины последовательности, что приводит к значительному улучшению как коммуникационного, так и вычислительного параллелизма, а также их перекрытия. Кроме того, мы расширяем LASP-2 до LASP-2H, применяя аналогичную переработку коммуникации к стандартным модулям внимания, предлагая эффективное решение SP для гибридных моделей, объединяющих линейные и стандартные слои внимания. Наша оценка на модели Linear-Llama3, варианте Llama3 с линейным вниманием вместо стандартного внимания, демонстрирует эффективность LASP-2 и LASP-2H. В частности, LASP-2 достигает улучшения скорости обучения на 15,2% по сравнению с LASP и на 36,6% по сравнению с Ring Attention при длине последовательности 2048K на 64 GPU. Код выпущен как часть: https://github.com/OpenSparseLLMs/Linear-MoE.

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.