Lightning Attention-2: Una soluzione gratuita per gestire lunghezze di sequenza illimitate nei grandi modelli linguistici

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

January 9, 2024

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

cs.AI

Abstract

L'attenzione lineare è un meccanismo di attenzione efficiente che è emerso recentemente come una valida alternativa all'attenzione softmax convenzionale. Grazie alla sua capacità di elaborare token con complessità computazionali lineari, l'attenzione lineare, in teoria, può gestire sequenze di lunghezza illimitata senza sacrificare la velocità, mantenendo cioè una velocità di addestramento costante per varie lunghezze di sequenza con un consumo di memoria fisso. Tuttavia, a causa del problema della somma cumulativa (cumsum), gli attuali algoritmi di attenzione lineare non riescono a dimostrare il loro vantaggio teorico in un contesto causale. In questo articolo, presentiamo Lightning Attention-2, la prima implementazione di attenzione lineare che consente all'attenzione lineare di realizzare i suoi benefici computazionali teorici. Per raggiungere questo obiettivo, sfruttiamo l'idea del tiling, gestendo separatamente i componenti intra-blocco e inter-blocco nel calcolo dell'attenzione lineare. Nello specifico, utilizziamo il meccanismo di calcolo dell'attenzione convenzionale per gli intra-blocchi e applichiamo i trucchi del kernel di attenzione lineare per gli inter-blocchi. Una tecnica di tiling viene adottata sia nelle procedure di forward che di backward per sfruttare al massimo l'hardware GPU. Implementiamo il nostro algoritmo in Triton per renderlo IO-aware e hardware-friendly. Sono stati condotti vari esperimenti su diverse dimensioni di modelli e lunghezze di sequenza. Lightning Attention-2 mantiene una velocità di addestramento e inferenza costante indipendentemente dalla lunghezza della sequenza di input ed è significativamente più veloce rispetto ad altri meccanismi di attenzione. Il codice sorgente è disponibile all'indirizzo 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.