Lightning Attention-2: Een Gratis Oplossing voor het Omgaan met Onbeperkte Sequentielengtes in Grote Taalmodellen

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

Samenvatting

Lineaire aandacht is een efficiënt aandachtmechanisme dat recentelijk naar voren is gekomen als een veelbelovend alternatief voor conventionele softmax-aandacht. Met zijn vermogen om tokens te verwerken in lineaire computationele complexiteiten, kan lineaire aandacht in theorie sequenties van onbeperkte lengte aan zonder snelheid in te leveren, d.w.z. een constante trainingssnelheid behouden voor verschillende sequentielengtes met een vast geheugengebruik. Echter, vanwege het probleem met cumulatieve sommatie (cumsum), kunnen huidige lineaire aandachtalgoritmen hun theoretische voordeel niet aantonen in een causale setting. In dit artikel presenteren we Lightning Attention-2, de eerste implementatie van lineaire aandacht die het mogelijk maakt om de theoretische computationele voordelen te realiseren. Om dit te bereiken, maken we gebruik van de gedachte van tiling, waarbij de intra-blok en inter-blok componenten in de lineaire aandachtberekening afzonderlijk worden behandeld. Specifiek gebruiken we het conventionele aandachtberekeningsmechanisme voor de intra-blokken en passen we lineaire aandachtkernel trucs toe voor de inter-blokken. Een tiling-techniek wordt toegepast in zowel de voorwaartse als de achterwaartse procedures om volledig gebruik te maken van de GPU-hardware. We implementeren ons algoritme in Triton om het IO-bewust en hardwarevriendelijk te maken. Diverse experimenten worden uitgevoerd op verschillende modelgroottes en sequentielengtes. Lightning Attention-2 behoudt een consistente trainings- en inferentiesnelheid ongeacht de invoersequentielengte en is aanzienlijk sneller dan andere aandachtmechanismen. De broncode is beschikbaar op 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.