Fundamentele AI

Cache-ul dvs. KV nu are o problemă de biți. Are o problemă de geometrie.

mm
Adaugă Unite.AI la sursele tale preferate pe Google

La aceeași precizie de 2 biți, o decizie despre care axă se cuantizează poate schimba scorul benchmark de la 2,88 la 63,53. Cheile și valorile necesită un tratament opus – și motivul se află în ecuația de atenție, nu în hardware.

Luați Llama-2-13B. Grupați cache-ul cheie-valoare după o dimensiune de grup de cuantizare de 32 în două biți, lăsând totul altceva în loc – același model, același buget de biți, aceleași dimensiuni de grup, aceleași benchmark-uri.

În funcție de un singur element al unei decizii de implementare, rezultatele acurateței CoQA pot fi fie 2,88, fie 63,53. Scorul folosind precizia completă este 66,37.

Decizia nu este despre câți biți sunt folosiți în total. Întrebarea este pur și simplu ce axă alegeți să grupați atunci când calculați fiecare factor de scară? Când decideți să folosiți canalul ca dimensiune de grupare (chei) și tokenul ca dimensiune de grupare (valori), ajungeți undeva în jurul a patru puncte de la performanța preciziei complete. Dacă schimbați una dintre aceste alegeri, experimentați o pierdere de calitate. Dacă schimbați ambele alegeri, modelul nu mai funcționează.

Patru moduri de a cheltui aceiași 2 biți pe același cache. Rezultate din ablația KIVI pe Llama-2-13B la dimensiunea grupului 32.

Cuantizarea este de obicei considerată ca fiind doar un singur regulator: 8 biți, 4 biți, 2 biți, cu un cost de precizie netedă. În interiorul cache-ului KV, nu este așa. Este vorba despre alegerea sistemelor de coordonate și despre faptul că sistemele diferite se aplică cheilor și valorilor. Acest articol explică de ce. Pe scurt: eroarea de cuantizare depinde de intervalul de valori din grupuri; cheile și valorile au structuri foarte diferite; și oamenii adesea greșesc pentru că nu puteți deriva axa corectă din distribuția valorilor deloc. Trebuie să vă uitați la modul în care eroarea se schimbă după ce atenția o consumă. Acest lucru oferă un principiu general pentru comprimarea activărilor intermediare și un motiv bun pentru a pune la îndoială eroarea de reconstrucție ca proxy pentru calitate.

De ce cache-ul KV este locul în care se întâmplă acest lucru

În timpul fazei de generare, un transformator stochează toate datele de proiecție cheie-valoare (KV) ale tokenilor pe care i-a procesat anterior într-un cache, astfel încât să nu fie nevoie să recalculeze aceste date din nou. Acest cache crește liniar cu lungimea contextului și dimensiunea lotului. În cele din urmă, acest lucru va duce la faptul că cache-ul va crește mai mare decât modelul însuși.

Această creștere a dimensiunii poate fi ușor identificată atunci când se examinează consumul de memorie al diferitelor părți ale modelului. În analiza KVQuant a LLaMA-7B, greutățile reprezintă aproximativ 98 la sută din memorie la o lungime de secvență de 512, cu activări la 2 la sută. La 128K context, raportul se inversează la aproximativ 16 la sută greutăți și 84 la sută cache KV. Când examinăm o analiză a OPT-175B citată de autorii KIVI, ei au găsit rezultate similare. În special, la o dimensiune de lot de 512 cu un prompt de 512 tokeni, cache-ul KV ajunge la 1,2 TB – de mai multe ori dimensiunea greutăților modelului.

Cu toate acestea, capacitatea este doar jumătate din problema de aici. GPU-ul trebuie să citească întregul cache KV din memoria dispozitivului pentru fiecare token generat. Acest lucru înseamnă că, în timp ce GPU-ul citește cache-ul KV, nucleele de calcul stau inactiv. Prin urmare, reducerea dimensiunii totale a cache-ului atât crește spațiul de procesare disponibil, cât și reduce timpul petrecut așteptând transferurile de date.

Ce este de fapt eroarea de cuantizare

Cuantizarea integer uniformă este matematică simplă. Pentru un grup de numere, se înregistrează cel mai mic număr ca punct zero și apoi se împarte intervalul grupului la numărul de niveluri care pot fi reprezentate pentru a obține o dimensiune a pasului. Apoi, se rotunjește fiecare element la pasul cel mai apropiat. Două rezultate imediate urmează. Prima, eroarea pe element este limitată la jumătate de pas. A doua, dimensiunea pasului este intervalul grupului împărțit la 2^B – 1. La 2 biți, aveți doar 4 niveluri pentru a acoperi orice interval care există în acel grup. Deci, un element care este de o sută de ori mai mare în comparație cu vecinii săi nu se desfășoară doar prost. Inflaționează dimensiunea pasului pentru toți ceilalți elemente care împărtășesc același grup și toate devin mai grosolane împreună. Grupul este unitatea de deteriorare. Alegerea unei axe înseamnă a decide care elemente suferă împreună. Punând întrebarea într-un mod diferit, nu mai este “câți biți pot să-mi permit să sacrific?” ci “unde sunt valorile extreme și pot să le izolez?”

Cheile: outlier-ii trăiesc în canale fixe

Modelele de limbaj mare conțin activări care sunt neobișnuit de mari în comparație cu majoritatea activărilor. Sun și colegii săi au catalogat aceste activări foarte mari în diferite familii de modele: în Mixtral 8x7B, magnitudinea cea mai mare este aproape de 7000, în timp ce magnitudinea medie a caracteristicii este de aproximativ 0,3 – aproximativ patru ordine de mărime diferite. Acestea sunt foarte rare; ele rămân fixe în dimensiuni care rareori se schimbă cu intrarea și nu sunt accidentale. Ele acționează ca biasuri implicite și sunt ceea ce focalizează atenția asupra doar câtorva tokeni: comportamentul de scurgere a atenției. În cache-ul cheie, această structură este foarte clară: anumite canale transportă magnitudini foarte mari consistent de-a lungul fiecărui token dintr-o secvență. Grupați după tokeni și fiecare grup conține acele canale outlier, astfel încât dimensiunea pasului pentru fiecare grup este setată de canalele outlier și toate canalele obișnuite plătesc pentru asta. Grupați după canale și canalele outlier formează propriile grupuri. Intervalul lor intern este mare, dar autoconținut; canalele obișnuite sunt lăsate singure. Rezultatele se potrivesc. Mediate pe straturi și capete pe Llama-2-13B, KIVI raportează o eroare de reconstrucție a cheii de 13,67 sub gruparea pe token, față de 4,55 pe canal, și – mai important – o eroare de scor de atenție de 47,00 față de 9,60. Cuantizarea cheilor pe token produce aproximativ de cinci ori mai multă eroare de scor. Scorurile se potrivesc cu metricile semnificative pentru chei; cuantizarea pe canal excelează pe ambele fronturi.

Valorile: unde se întrerupe intuiția

Cache-ul valorilor nu prezintă un model de outlier de canal. Pare a fi destul de plat. Pe cont propriu, prin argumentul intervalului, am putea să ne așteptăm ca oricare dintre aceste axe să producă o calitate similară de compresie.

Ele nu o fac. Indiferent de modul în care se implementează gestionarea cheilor (rezultatele de 2,80 și 2,88), comprimarea valorilor per-canal colapsează modelul.

Și aici este capcana: dacă măsurați această pierdere folosind eroarea de reconstrucție brută pe tensorul original pentru care fiecare valoare a fost comprimată, cuantizarea valorilor per-canal pare de fapt puțin mai bună, cu 3,73 față de 4,57. Dacă ați valida comprimarea în mod evident, ați alege configurația care distruge modelul.

Eroarea de cuantizare a cache-ului valorilor pe Llama-2-13B, măsurată în două moduri. Metrica tensorului stocat și metrica ieșirii consumate diferă cu mai mult de o ordine de mărime.

Rezoluția constă în faptul că cache-ul valorilor nu este citit niciodată direct. Este consumat de un produs matricial: ieșirea atenției este o sumă ponderată a vectorilor de valori de-a lungul tokenilor, cu scoruri de atenție softmax ca ponderi. Din cauza acestui lucru, eroarea relevantă este cea introdusă în timpul acestui proces și nu în interiorul tensorilor înșiși. Măsurată în funcție de ieșirea atenției, ordinea a fost complet inversată. Eroarea relativă raportată de KIVI pentru ieșirea atenției datorită cuantizării vectorului de valori pe token a fost de 3,55 față de 49,89 pentru cuantizarea pe canal – de peste paisprezece ori mai mare pentru ceea ce părea alegerea mai bună pe baza modului în care a fost comprimată.

Explicația constă în raritatea atenției, pe care au măsurat-o ca fiind de 84,3 la sută. Majoritatea informațiilor conținute în ieșire pot fi atribuite unui număr mic de tokeni foarte importanți. Cuantizarea pe token limitează eroarea fiecărui token la acel token, astfel încât erorile pe tokeni neimportanți se multiplică cu ponderi de atenție aproape de zero și efectiv dispar. Cuantizarea pe canal răspândește eroarea fiecărui token pe un canal comun, astfel încât tokenii prost reprezentați contaminează reprezentarea celor care contează. Raritatea care face ca atenția să fie eficientă este aceeași proprietate care face ca cuantizarea pe token să fie sigură.

Lecția transferabilă este mai amplă decât cache-ul KV: măsurați eroarea de compresie acolo unde tensorul este consumat, nu acolo unde este stocat. O presupunere implicită făcută de eroarea de reconstrucție este că fiecare componentă a unui tensor are o greutate egală atunci când contribuie la ieșirea finală. Atenția nu o face în mod explicit. Orice operație downstream care cântărește, pornește sau sparsifică intrarea sa rupe această presupunere. Citiitorii familiarizați cu articolul meu anterior despre orbirea contradicției în sistemele de recuperare vor recunoaște că aceste rezultate sunt similare cu eșecurile descrise anterior: metrici ușor de calculat care raportează altceva decât ceea ce a fost intenționat.

Încorporările rotative complicate cheile

Există unele probleme cu utilizarea încorporărilor rotative de poziție (RoPE). RoPE rotite perechi de canale în funcție de poziția relativă a fiecărui token. Această amestecare parțială dizolvă structura de canal fixă care a făcut cuantizarea cheilor per-canal să funcționeze în primul rând – un canal outlier se rotitește în vecinii săi și vecinii moștenesc intervalul. Răspunsul KVQuant este ordinea: cuantizați cheile înainte de a aplica rotația și aplicați RoPE după de-cuantizare. Împreună cu cuantizarea cheilor per-canal, tipuri de date neuniforme și izolarea unei fracțiuni mici de outlieri, aceasta obține sub 0,1 degradare de perplexitate la 3 biți și permite servirea LLaMA-7B până la 1 milion de tokeni de context pe un singur A100-80GB.

De asemenea, este important să înțelegem nivelul de impact din RoPE. Autorii lucrării “RotateKV” au raportat o creștere de 145 la sută a erorilor de cuantizare odată ce RoPE a fost adăugat și au observat că canalele outlier diferă între capetele de atenție – ceea ce face ca aplicarea unei matrice de rotație partajate peste tot să fie insuficientă și rotațiile adaptate capetelor să facă mai bine.

Impozitul sistemelor și de ce nu este un detaliu

Cuantizarea pe token se potrivește bine decodării. Fiecare token sosește; îl cuantizați, îl adăugați la secvență (de-a lungul dimensiunii tokenului), nimic altceva nu se mișcă.

Cu toate acestea, cuantizarea pe canal nu se potrivește. Deoarece statisticile unui canal se întind pe tokeni care nu au fost încă generați, nu puteți calcula un factor de scară atunci când un token intră. Soluția KIVI este să păstrați tokenii cei mai recent generați – până la 128 – în precizie completă într-un tampon rezidual și să cuantizați în grupuri odată ce s-au acumulat suficienți.

Așa se întâmplă că tamponul rezidual devine suport, mai degrabă decât doar o chestiune incidentală. Pe GSM8K cu Llama-2-7B, scorurile de precizie completă sunt de 13,50. Cuantizate complet la 2 biți cu axe corecte, scorurile sunt de 5,76. Aceleași axe și aceiași biți, plus tamponul rezidual de tokeni recent generați în precizie completă, scorurile sunt de 12,74. O fereastră de tokeni recent generați în precizie completă va recupera mult din ceea ce a fost pierdut din cauza cuantizării agresive pe probleme dificile – ceea ce ar fi logic dacă ne gândim la care tokeni erau atenți de o serie de operații aritmetice.

Există un beneficiu semnificativ în a face toate aceste lucruri corect – așa cum raportează KIVI, 2,6 ori mai puțină utilizare a memoriei pentru Llama-2-7B, permițând dimensiuni de lot de până la 4 ori mai mari, precum și 2,35 până la 3,47 ori mai bine flux de lucru pe o sarcină de serviciu din lumea reală.

Ce să faceți cu asta

  1. Nu folosiți un singur cuantizator pentru ambele. Folosiți cuantizatori diferiți pentru chei (pe canal) și pentru valori (pe token). O pipeline care aplică un singur cuantizator la “cache-ul KV” a sacrificat probabil deja cea mai mare parte a calității posibile atunci când se folosesc câțiva biți pentru a reprezenta fiecare valoare.
  2. Cuantizați cheile înainte de RoPE. Acesta este un aspect al corectitudinii, mai degrabă decât o chestiune de preferință.
  3. Păstrați o fereastră de tokeni recent generați în precizie completă. Deși păstrarea unei astfel de ferestre necesită foarte puțină memorie în comparație cu cât de mare poate fi un cache, este exact această zonă care generează multă acuratețe pentru sarcini dificile.
  4. Nu validați pe eroarea de reconstrucție. Validați întotdeauna pe baza ieșirii atenției sau pe baza performanței sarcinii finale. Metrica de stocare nu este doar zgomotoasă – pentru valori, ea indică în direcția greșită.
  5. Nu validați pe benchmark-uri de alegere multiplă pe context scurt. Autorii KIVI evită în mod deliberat sarcini închise, cum ar fi MMLU, pentru această evaluare, deoarece un singur pas de decodare care citește logit-urile de ieșire abia exercită cache-ul. Orice evaluare care nu construiește un cache în timp și apoi efectuează generare din el nu va putea niciodată să observe eșecurile inerente din proiectarea sistemului.

Încotro se îndreaptă lucrul

Deși mai există încă ceva de făcut cu privire la natura geometrică a problemei, mulți cercetători continuă să studieze modul în care canalele outlier sunt distribuite între diferitele capete de atenție ale transformatorului și modul în care limitările hardware afectează care grupări sunt mai ieftine: InnerQ împătură normalizarea cheilor canalului în greutățile cheii și interogării în timpul preumplerii. Prin urmare, nu se incorporează nici o suprasarcină la runtime. Mai mult, InnerQ stochează ferestre de precizie ridicată atât pentru tokenii recent generați, cât și pentru tokenii care primesc atenție. Făcând acest lucru, InnerQ elimină oportunitatea ca outlierii din canalul care primește atenție să contamineze canalele vecine.

Alții propun că, în loc să stocheze întregul cache, ar trebui să stocheze doar suficiente informații pentru a putea reface cheia și/sau valoarea (valorile) la cerere dintr-o reprezentare mai mică a cache-ului.

În cele din urmă, este important să rețineți că precizia nu este singurul parametru pe care cuantizarea îl afectează. Cercetarea publicată recent a demonstrat degradarea alinierii rezultate din cuantizarea cache-ului KV. Mai mult, această cercetare a documentat degradarea alinierii chiar și în medii de servicii vLLM de producție care utilizează cache-uri FP8, împreună cu un protocol de recuperare fără antrenare care a restabilit până la 97 la sută din ceea ce a fost pierdut în ceea ce privește alinierea. Prin urmare, în timp ce o configurație poate să-și păstreze rezultatele benchmark, nu înseamnă neapărat că reține toți ceilalți parametri pe care îi îngrijiți.

Principiul general

Ideea de cuantizare a fost înfățișată ca un “buget de precizie”: câți biți pot să-mi permit să sacrific? Cache-ul KV arată că întrebarea mai utilă este structurală. Precizia este alocată în grupuri; grupul este unitatea de deteriorare, iar axa pe care grupați determină care elemente împărtășesc soarta lor. Axa corectă este cea pe care tensorul este consumat, adică modul în care utilizați tensorul și nu modul în care tensorul apare atunci când este stocat în memorie. Cheile sunt utilizate prin intermediul unui calcul de produs punctual împotriva interogării. Un singur canal corupt va otrăvi toate scorurile. Valorile sunt consumate printr-o computație de medie ponderată pe tokeni, cu scoruri de atenție softmax ca ponderi. Prin urmare, un token corupt este pur și simplu ponderat.

Două tensore de dimensiuni identice și generate de straturi consecutive sunt tratate diferit. Este important să întrebați orice activare pe care intenționați să o comprimați: ce operație contractează asta și respectă gruparea mea?

Himanshu Goel este un cercetător AI/ML specializat în generarea augmentată de recuperare pentru domenii cu risc ridicat, incluzând fluxuri de lucru biomedicale, financiare și de reglementare a documentelor.