Hur du snabbar upp dina AI-modeller – utan att steka ditt minne

Tabell över länkar

Abstrakt och 1 inledning

2. Relaterat arbete

3. Bakgrund

3.1. Notation och 3.2. Språkmodellinferens

3.3. Multi-Query, Multi-Head och Generalized Multi-Query Attention

4. Kontextmedveten bifurkerad uppmärksamhet och 4.1. Motivering

4.2. Formulering och 4.3. Minnes IO-komplexitet

5. Experiment

5.1. Jämföra funktionerna för Multi-Head, Multi-Query och Multi-Group Attention

5.2. Latenser för kapacitetslikvärdiga modeller

5.3. Ansökningar

6. Slutsats och referenser

A. Vanliga frågor

B. Relaterat arbete

C. Inställning

D. Multi-Group Attention Family

E. Kontextmedveten bifurkerad uppmärksamhet

F. Ansökningar: Ytterligare resultat

G. Kompatibilitet med tekniker för spekulativ avkodning och snabb avkodning

Abstrakt

I vår studie presenterar vi bifurcated attention , en metod som utvecklats för språkmodellinferens i enkontext-batchsamplingssammanhang. Detta tillvägagångssätt syftar till att minska redundanta minnes-IO-kostnader, en betydande faktor för latens för höga batchstorlekar och långa sammanhangslängder. Bifurcated uppmärksamhet uppnår detta genom att dela upp uppmärksamhetsmekanismen under inkrementell avkodning i två distinkta GEMM-operationer, med fokus på KV-cachen från förfyllning och avkodningsprocessen. Denna metod säkerställer exakt beräkning och upprätthåller den vanliga beräkningsbelastningen (FLOPs) för standarduppmärksamhetsmekanismer, men med reducerad minnes-IO. Bifurcated uppmärksamhet är också kompatibel med multi-query uppmärksamhetsmekanism känd för reducerad minnes-IO för KV-cache, vilket ytterligare möjliggör högre batchstorlek och kontextlängd. Den resulterande effektiviteten leder till lägre latens, vilket förbättrar lämpligheten för realtidsapplikationer, t.ex. möjliggör massivt parallella svarsgenerering utan att avsevärt öka latensen, vilket förbättrar prestandan när den integreras med efterbehandlingstekniker såsom omrangering.

1. Inledning

Tillkomsten av stora språkmodeller (LLM) har inlett en ny era av maskininlärning, som uppvisar anmärkningsvärda prestanda på ett brett spektrum av uppgifter (Brown et al., 2020; OpenAI, 2023; Chowdhery et al., 2022; Touvron et al., 2023; et al., 2021.; Hoffmann et al., 2021.; Chen et al., 2021.; Li et al., 2022; Amazon, 2022; Trots deras imponerande kapacitet innebär utplaceringen av dessa storskaliga modeller i praktiska tillämpningar betydande utmaningar, särskilt när det gäller slutledningslatens och effektivitet. Att förbättra dessa aspekter är avgörande, eftersom de direkt påverkar de beräkningsresurser som krävs för att generera förutsägelser och möjliggöra praktisk implementering av dessa avancerade modeller inom olika branscher.

Ett särskilt krävande slutledningsscenario är batchsampling med en enda kontext, där målet är att generera flera kompletteringar från ett enda sammanhang. Den här uppgiften förekommer ofta i många applikationer som kodredigering av IDE-verktyg som ger flera rekommendationer, eller i fall där rankning bland många generationer behövs för optimal prestanda (via rankningsmått som genomsnittlig loggsannolikhet, majoritetsomröstning, etc). Den inkrementella avkodningen av ett sådant samplingsscenario är minnes-IO-intensiv, vilket blir en latensflaskhals för höga batcher och kontextlängder.

I den här studien undersöker vi två kompatibla strategier för att ta itu med minnes-IO-utmaningarna i transformatorinferens: (1) en undersökning av multi-query och dess avvägningar, och (2) en ny teknik som kallas kontextmedveten bifurkerad uppmärksamhet.

Vår undersökning börjar med en analys av den generaliserade multi-query uppmärksamheten (Ainslie et al., 2023), som inkluderar multi-query (Shazeer, 2019), såväl som den etablerade multi-head uppmärksamhet mekanismen (Vaswani et al., 2017) för prestanda och latens trade-off. Våra resultat visar smidig prestandaskalning med ökande modellstorlek för ett fast värde på antalet grupper g för generaliserad multi-query[1]. En sänkning av g resulterar i en förskjutning uppåt av skalningskurvorna för valideringsförlusten mot modellstorleken. Det konsekventa förhållandet mellan cachekomprimeringen, modellstorleken och valideringsförlusten gör att vi kan avväga slutledningseffektivitet med modellstorlek, dvs. gör det möjligt för oss att välja högre komprimering för användningsfall som kräver hög effektivitet, samtidigt som vi matchar prestandan för multi-head uppmärksamhet genom att kompensera med en större modellstorlek.

För det andra introducerar vi kontextmedveten bifurkerad uppmärksamhet, en teknik som delar all uppmärksamhet i den generaliserade flerfrågefamiljen till kontext- och avkodningskomponenter under inkrementell avkodning. Sådan bifurkation involverar samma antal FLOP:ar och ger identiska resultat jämfört med den ursprungliga uppmärksamheten, men kan avsevärt minska minnes-IO-kostnaden och därmed latensen i scenarier med hög batch- och kontextlängd. Detta tillvägagångssätt tillåter generering av flera realtidsavslut utan att dra på sig mycket extra latenskostnader, eller möjliggör mycket högre batchstorlekar vilket leder till förbättrad rankningsprestanda. Till exempel, för CodeGen 16B multi-head-modell (Nijkamp et al., 2022) med 2k kontextlängd, kan vi öka batchstorleken till 128 med bifurcerad uppmärksamhet, jämfört med batchstorlek på endast 5 utan, vilket resulterar i att pass@k (Chen et al., 2021) ökar från 59,06% till 59,06%, till 59,06% via medelvärde. 55,2 % till 58,1 %.

[1] Lägre värden för uppmärksamhetsgrupper g leder till högre komprimering av nyckel-värde-tensorerna, som i multi-query-fallet där g = 1, vilket förbättrar slutledningseffektiviteten och latensen på grund av minskad KV-cache jämfört med multi-head-fallet där g = h, antalet frågeuppmärksamhetshuvuden.