Lightning Attention-2: Um Almoço Grátis para Lidar com Comprimentos de Sequência Ilimitados em Modelos de Linguagem de Grande Escala

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

January 9, 2024

Autores: Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, Yiran Zhong

cs.AI

Resumo

A atenção linear é um mecanismo de atenção eficiente que recentemente surgiu como uma alternativa promissora à atenção softmax convencional. Com sua capacidade de processar tokens em complexidades computacionais lineares, a atenção linear, em teoria, pode lidar com sequências de comprimento ilimitado sem sacrificar a velocidade, ou seja, mantendo uma velocidade de treinamento constante para vários comprimentos de sequência com um consumo de memória fixo. No entanto, devido ao problema com a soma cumulativa (cumsum), os algoritmos atuais de atenção linear não conseguem demonstrar sua vantagem teórica em um cenário causal. Neste artigo, apresentamos o Lightning Attention-2, a primeira implementação de atenção linear que permite que a atenção linear realize seus benefícios computacionais teóricos. Para alcançar isso, aproveitamos a ideia de tiling, tratando separadamente os componentes intra-bloco e inter-bloco no cálculo da atenção linear. Especificamente, utilizamos o mecanismo convencional de computação de atenção para os intra-blocos e aplicamos truques de kernel de atenção linear para os inter-blocos. Uma técnica de tiling é adotada tanto nos procedimentos de forward quanto de backward para aproveitar ao máximo o hardware da GPU. Implementamos nosso algoritmo em Triton para torná-lo consciente de IO e amigável ao hardware. Vários experimentos são conduzidos em diferentes tamanhos de modelo e comprimentos de sequência. O Lightning Attention-2 mantém uma velocidade de treinamento e inferência consistente, independentemente do comprimento da sequência de entrada, e é significativamente mais rápido do que outros mecanismos de atenção. O código-fonte está disponível em 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.