Lightning Attention-2: Бесплатный обед для обработки неограниченной длины последовательностей в больших языковых моделях

Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

January 9, 2024

Авторы: Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, Yiran Zhong

cs.AI

Аннотация

Линейное внимание — это эффективный механизм внимания, который недавно появился как перспективная альтернатива традиционному softmax-вниманию. Благодаря способности обрабатывать токены с линейной вычислительной сложностью, линейное внимание теоретически может работать с последовательностями неограниченной длины без потери скорости, то есть поддерживать постоянную скорость обучения для различных длин последовательностей при фиксированном потреблении памяти. Однако из-за проблемы с кумулятивным суммированием (cumsum) текущие алгоритмы линейного внимания не могут продемонстрировать свои теоретические преимущества в причинно-следственных настройках. В данной статье мы представляем Lightning Attention-2 — первую реализацию линейного внимания, которая позволяет линейному вниманию реализовать свои теоретические вычислительные преимущества. Для этого мы используем идею тайлинга, раздельно обрабатывая внутриблочные и межблочные компоненты в вычислении линейного внимания. В частности, мы применяем традиционный механизм вычисления внимания для внутриблочных компонентов и используем трюки с ядрами линейного внимания для межблочных компонентов. Техника тайлинга применяется как в прямом, так и в обратном проходе, чтобы максимально эффективно использовать аппаратные возможности GPU. Мы реализуем наш алгоритм в Triton, чтобы сделать его IO-оптимизированным и дружественным к оборудованию. Проведены различные эксперименты на моделях разных размеров и последовательностях разной длины. Lightning Attention-2 сохраняет постоянную скорость обучения и вывода независимо от длины входной последовательности и значительно быстрее других механизмов внимания. Исходный код доступен по адресу https://github.com/OpenNLPLab/lightning-attention.

English

Linear attention is an efficient attention mechanism that has recently emerged as a promising alternative to conventional softmax attention. With its ability to process tokens in linear computational complexities, linear attention, in theory, can handle sequences of unlimited length without sacrificing speed, i.e., maintaining a constant training speed for various sequence lengths with a fixed memory consumption. However, due to the issue with cumulative summation (cumsum), current linear attention algorithms cannot demonstrate their theoretical advantage in a causal setting. In this paper, we present Lightning Attention-2, the first linear attention implementation that enables linear attention to realize its theoretical computational benefits. To achieve this, we leverage the thought of tiling, separately handling the intra-block and inter-block components in linear attention calculation. Specifically, we utilize the conventional attention computation mechanism for the intra-blocks and apply linear attention kernel tricks for the inter-blocks. A tiling technique is adopted through both forward and backward procedures to take full advantage of the GPU hardware. We implement our algorithm in Triton to make it IO-aware and hardware-friendly. Various experiments are conducted on different model sizes and sequence lengths. Lightning Attention-2 retains consistent training and inference speed regardless of input sequence length and is significantly faster than other attention mechanisms. The source code is available at https://github.com/OpenNLPLab/lightning-attention.