Lightning Attention-2: Ein kostenloser Ansatz zur Bewältigung unbegrenzter Sequenzlängen in großen Sprachmodellen

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

January 9, 2024

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

cs.AI

Zusammenfassung

Lineare Aufmerksamkeit ist ein effizienter Aufmerksamkeitsmechanismus, der sich kürzlich als vielversprechende Alternative zur konventionellen Softmax-Aufmerksamkeit herausgestellt hat. Mit der Fähigkeit, Token in linearer Rechenkomplexität zu verarbeiten, kann lineare Aufmerksamkeit theoretisch Sequenzen unbegrenzter Länge verarbeiten, ohne die Geschwindigkeit zu opfern, d. h. sie hält eine konstante Trainingsgeschwindigkeit für verschiedene Sequenzlängen bei festem Speicherverbrauch aufrecht. Aufgrund des Problems mit der kumulativen Summation (cumsum) können aktuelle lineare Aufmerksamkeitsalgorithmen jedoch ihren theoretischen Vorteil in einem kausalen Setting nicht demonstrieren. In diesem Artikel präsentieren wir Lightning Attention-2, die erste Implementierung linearer Aufmerksamkeit, die es ermöglicht, die theoretischen Rechenvorteile der linearen Aufmerksamkeit zu realisieren. Um dies zu erreichen, nutzen wir den Gedanken des Tiling, bei dem die intra-block und inter-block Komponenten in der Berechnung der linearen Aufmerksamkeit separat behandelt werden. Insbesondere verwenden wir den konventionellen Aufmerksamkeitsberechnungsmechanismus für die intra-blocks und wenden lineare Aufmerksamkeits-Kernel-Tricks für die inter-blocks an. Eine Tiling-Technik wird sowohl im Vorwärts- als auch im Rückwärtsverfahren eingesetzt, um die Vorteile der GPU-Hardware voll auszuschöpfen. Wir implementieren unseren Algorithmus in Triton, um ihn IO-bewusst und hardwarefreundlich zu gestalten. Verschiedene Experimente werden mit unterschiedlichen Modellgrößen und Sequenzlängen durchgeführt. Lightning Attention-2 behält eine konsistente Trainings- und Inferenzgeschwindigkeit unabhängig von der Eingabesequenzlänge bei und ist deutlich schneller als andere Aufmerksamkeitsmechanismen. Der Quellcode ist verfügbar unter 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.