Lightning Attention-2: Una solución sin coste adicional para manejar longitudes de secuencia ilimitadas en modelos de lenguaje grandes

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

Resumen

La atención lineal es un mecanismo de atención eficiente que ha surgido recientemente como una alternativa prometedora a la atención softmax convencional. Con su capacidad para procesar tokens en complejidades computacionales lineales, la atención lineal, en teoría, puede manejar secuencias de longitud ilimitada sin sacrificar la velocidad, es decir, manteniendo una velocidad de entrenamiento constante para varias longitudes de secuencia con un consumo de memoria fijo. Sin embargo, debido al problema con la suma acumulativa (cumsum), los algoritmos actuales de atención lineal no pueden demostrar su ventaja teórica en un entorno causal. En este artículo, presentamos Lightning Attention-2, la primera implementación de atención lineal que permite a la atención lineal alcanzar sus beneficios computacionales teóricos. Para lograrlo, aprovechamos la idea de la técnica de teselado, manejando por separado los componentes intra-bloque e inter-bloque en el cálculo de la atención lineal. Específicamente, utilizamos el mecanismo de cálculo de atención convencional para los intra-bloques y aplicamos trucos de kernel de atención lineal para los inter-bloques. Se adopta una técnica de teselado tanto en los procedimientos de avance como de retroceso para aprovechar al máximo el hardware de la GPU. Implementamos nuestro algoritmo en Triton para hacerlo consciente de las operaciones de entrada/salida (IO) y amigable con el hardware. Se realizan varios experimentos en diferentes tamaños de modelos y longitudes de secuencia. Lightning Attention-2 mantiene una velocidad de entrenamiento e inferencia constante independientemente de la longitud de la secuencia de entrada y es significativamente más rápido que otros mecanismos de atención. El código fuente está disponible en 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.