Lightning Attention-2 : Une solution sans compromis pour gérer des séquences de longueur illimitée dans les grands modèles de langage

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

January 9, 2024

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

cs.AI

Résumé

L'attention linéaire est un mécanisme d'attention efficace qui a récemment émergé comme une alternative prometteuse à l'attention softmax conventionnelle. Grâce à sa capacité à traiter les tokens avec des complexités computationnelles linéaires, l'attention linéaire peut, en théorie, gérer des séquences de longueur illimitée sans sacrifier la vitesse, c'est-à-dire en maintenant une vitesse d'entraînement constante pour différentes longueurs de séquences avec une consommation de mémoire fixe. Cependant, en raison du problème lié à la sommation cumulative (cumsum), les algorithmes actuels d'attention linéaire ne peuvent pas démontrer leur avantage théorique dans un cadre causal. Dans cet article, nous présentons Lightning Attention-2, la première implémentation d'attention linéaire qui permet à l'attention linéaire de réaliser ses avantages computationnels théoriques. Pour y parvenir, nous exploitons l'idée de la division en tuiles, en traitant séparément les composants intra-blocs et inter-blocs dans le calcul de l'attention linéaire. Plus précisément, nous utilisons le mécanisme de calcul d'attention conventionnel pour les intra-blocs et appliquons des astuces de noyau d'attention linéaire pour les inter-blocs. Une technique de division en tuiles est adoptée à la fois dans les procédures avant et arrière pour tirer pleinement parti du matériel GPU. Nous implémentons notre algorithme dans Triton pour le rendre conscient des entrées-sorties (IO-aware) et adapté au matériel. Diverses expériences sont menées sur différentes tailles de modèles et longueurs de séquences. Lightning Attention-2 conserve une vitesse d'entraînement et d'inférence constante quelle que soit la longueur de la séquence d'entrée et est significativement plus rapide que les autres mécanismes d'attention. Le code source est disponible à l'adresse 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.