Il tuo cache KV non ha un problema di bit. Ha un problema di geometria.

A precisione di 2 bit identica, una decisione su quale asse quantizzare fa oscillare il punteggio di un benchmark da 2,88 a 63,53. Le chiavi e i valori richiedono un trattamento opposto – e il motivo è nell’equazione dell’attenzione, non nel hardware.

Prendiamo Llama-2-13B. Raggruppiamo la sua cache di chiavi e valori per una dimensione di gruppo di quantizzazione di 32 in due bit, lasciando tutto il resto invariato – stesso modello, stesso budget di bit, stesse dimensioni di gruppo, stessi benchmark.

A seconda di un singolo elemento di una decisione di implementazione, i risultati di accuratezza di CoQA possono essere di 2,88 o 63,53. Il punteggio utilizzando la precisione completa è 66,37.

La decisione non riguarda il numero totale di bit utilizzati. La domanda è semplicemente quale asse scegliere per raggruppare quando si calcola ogni fattore di scala? Quando si decide di utilizzare il canale come dimensione di raggruppamento (chiavi) e il token come dimensione di raggruppamento (valori), si finisce da qualche parte all’interno di quattro punti della prestazione di precisione completa. Se si capovolge una di queste scelte, si verifica una perdita di qualità. Se si capovolge entrambe queste scelte, il modello non funziona più.

Quattro modi per spendere gli stessi 2 bit sulla stessa cache. Risultati dall’ablation KIVI su Llama-2-13B con dimensione di gruppo 32.

La quantizzazione è solitamente considerata come un’unica manopola: 8 bit, 4 bit, 2 bit, con un costo di accuratezza legato. All’interno della cache KV, non è così. È una scelta di sistemi di coordinate, e sistemi diversi si applicano a chiavi e valori. Questo articolo spiega perché. In breve: l’errore di quantizzazione dipende dalla gamma di valori all’interno dei gruppi; le chiavi e i valori hanno strutture molto diverse; e le persone spesso inciampano perché non si può derivare l’asse giusto dalla distribuzione dei valori. Bisogna guardare come l’errore cambia dopo che l’attenzione lo ha consumato. Ciò fornisce un principio generale per la compressione delle attivazioni intermedie e una buona ragione per dubitare dell’errore di ricostruzione come proxy della qualità.

Perché la cache KV è dove questo morde

Durante la fase di generazione, un trasformatore memorizza tutti i dati di proiezione di chiavi e valori (KV) dei token che ha precedentemente elaborato in una cache in modo che non debba ricalcolare questi dati nuovamente. Quella cache cresce linearmente con la lunghezza del contesto e la dimensione del batch. Alla fine, ciò porterà a una cache più grande del modello stesso.

Questo aumento di crescita può essere facilmente identificato quando si guarda al consumo di memoria delle diverse parti del modello. Nell’analisi KVQuant di LLaMA-7B, i pesi rappresentano circa il 98 percento della memoria a una lunghezza di sequenza di 512, con attivazioni al 2 percento. A 128K di contesto, il rapporto si inverte a circa il 16 percento di pesi e l’84 percento di cache KV. Quando guardiamo un’analisi di OPT-175B citata dagli autori di KIVI, troviamo risultati simili. In particolare, a una dimensione del batch di 512 con un prompt di 512 token, la cache KV raggiunge 1,2 TB – diverse volte le dimensioni dei pesi del modello.

Tuttavia, la capacità è solo metà del problema qui. La GPU deve leggere l’intera cache KV dalla memoria del dispositivo per ogni singolo token che genera. Ciò significa che mentre la GPU legge la cache KV, i core di calcolo rimangono inattivi. Pertanto, ridurre le dimensioni complessive della cache aumenta sia la testa di elaborazione disponibile che il tempo trascorso in attesa di trasferimenti di dati.

Cosa è realmente costituito l’errore di quantizzazione

La quantizzazione di interi uniformi è matematicamente semplice. Per un gruppo di numeri, si registra il numero più piccolo come punto zero e poi si divide la gamma del gruppo per il numero di livelli che possono essere rappresentati per ottenere una dimensione del passo. Quindi si arrotonda ogni elemento al passo più vicino. Due risultati immediati seguono. Primo, l’errore per elemento è limitato a metà passo. Secondo, la dimensione del passo è la gamma del gruppo divisa per 2ᴮ – 1. A 2 bit, si hanno solo 4 livelli per coprire qualsiasi diffusione esista all’interno del gruppo. Quindi un elemento che è cento volte più grande rispetto ai suoi vicini non si comporta solo male. Gonfia la dimensione del passo per tutti gli altri elementi che condividono lo stesso gruppo, e tutti diventano più grossolani insieme. Il gruppo è l’unità di danno. Scegliere un asse significa decidere quali elementi soffrono insieme. Inquadrare la domanda in modo diverso, non è più “quanti bit posso permettermi?” ma “dove sono i valori estremi e posso isolarli?”

Chiavi: gli outlier vivono in canali fissi

I grandi modelli linguistici contengono attivazioni che sono insolitamente grandi rispetto alla maggior parte delle attivazioni. Sun e colleghi hanno catalogato queste grandi attivazioni in diverse famiglie di modelli: in Mixtral 8x7B, la maggiore entità è vicina a 7000 mentre la mediana della grandezza della funzione è intorno a 0,3 – circa quattro ordini di grandezza diversi. Questi sono molto rari; rimangono fissi in dimensioni che raramente cambiano con l’input, e non sono accidentali. Agiscono come pregiudizi impliciti e sono ciò che focalizza l’attenzione su pochi token: il comportamento dell’attenzione. Nella cache delle chiavi, questa struttura è molto chiara: canali specifici trasportano grandezze molto grandi in modo coerente in ogni token in una sequenza. Raggruppare lungo i token e ogni gruppo contiene quegli outlier di canali, quindi ogni gruppo ha una dimensione del passo impostata dagli outlier, e tutti i canali ordinari pagano per questo. Raggruppare lungo i canali e gli outlier di canali formano i propri gruppi. La loro gamma interna è grande ma autonoma; i canali ordinari sono lasciati soli. I risultati corrispondono. Mediati su livelli e teste su Llama-2-13B, KIVI riporta un errore di ricostruzione di chiavi di 13,67 con raggruppamento per token contro 4,55 per canale, e – più importante – un errore di punteggio di attenzione di 47,00 contro 9,60. La quantizzazione delle chiavi per token produce circa cinque volte l’errore di punteggio. I punteggi concordano con metriche significative per le chiavi; la quantizzazione del canale eccelle in entrambi i fronti.

Valori: dove l’intuizione si rompe

La cache dei valori non mostra un modello di outlier di canali. Sembra essere abbastanza piatto. Da solo, per l’argomento della gamma, potremmo aspettarci che entrambi gli assi producano una qualità di compressione simile.

Non lo fanno. Indipendentemente da come viene implementata la gestione delle chiavi (i risultati 2,80 e 2,88), la compressione per canale dei valori fa collassare il modello.

E qui c’è l’inghippo: se si misura questa perdita utilizzando l’errore di ricostruzione grezzo sul tensore originale per cui ogni valore è stato compresso, la quantizzazione del valore per canale sembra in realtà leggermente migliore, con 3,73 contro 4,57. Se si convalida la propria compressione nel modo ovvio, si sceglierebbe la configurazione che distrugge il modello.

Errore di quantizzazione della cache dei valori su Llama-2-13B, misurato in due modi. La metrica del tensore archiviato e la metrica dell’output consumato non concordano di più di un ordine di grandezza.

La risoluzione è che la cache dei valori non viene letta direttamente. Viene consumata da un prodotto matriciale: l’output dell’attenzione è una somma ponderata di vettori di valori attraverso i token, con punteggi di attenzione softmax come pesi. A causa di ciò, l’errore rilevante è quello introdotto durante questo processo e non all’interno dei tensori stessi. Misurato in termini di output dell’attenzione, l’ordine era completamente invertito. L’errore relativo segnalato da KIVI per l’output dell’attenzione a causa della quantizzazione del vettore di valori per token era 3,55 rispetto a 49,89 per la quantizzazione per canale – più di quattordici volte superiore per ciò che sembrava la scelta migliore in base a come era stato compresso.

La spiegazione è la scarsità dell’attenzione, che hanno misurato come l’84,3 percento. La maggior parte delle informazioni contenute nell’output può essere attribuita a un piccolo numero di token molto importanti. La quantizzazione per token confina l’errore di ogni token a quel token, quindi gli errori sui token non importanti vengono moltiplicati per pesi di attenzione vicini a zero e di fatto scompaiono. La quantizzazione per canale sparge l’errore di ogni token attraverso una scala di canale condivisa, quindi i token male rappresentati contaminano la rappresentazione di quelli che contano. La scarsità che rende l’attenzione efficiente è la stessa proprietà che rende la quantizzazione per token sicura.

La lezione trasferibile è più ampia della cache KV: misurare l’errore di compressione dove il tensore viene consumato, non dove viene archiviato. Un’ipotesi implicita fatta dall’errore di ricostruzione è che ogni componente di un tensore abbia peso uguale quando contribuisce all’output finale. L’attenzione esplicitamente non lo fa. Qualsiasi operazione downstream che pesa, blocca o sparisce l’input rompe quell’ipotesi. I lettori familiari con il mio articolo precedente riguardante i punti ciechi nelle metriche di valutazione nei sistemi di recupero riconosceranno che questi risultati sono simili ai fallimenti precedentemente descritti: metriche facilmente calcolabili che segnalano qualcosa di diverso da ciò che era inteso.

Le incrustazioni rotarie complicano le chiavi

Ci sono alcuni problemi nell’utilizzo delle incrustazioni di posizione rotarie (RoPE). RoPE ruota le coppie di canali in base alla posizione relativa di ogni token. Quel misto dissolve parzialmente la struttura di canale fisso che ha reso la quantizzazione delle chiavi per canale funzionare nel primo posto – un canale outlier viene ruotato nei suoi vicini e i vicini ereditano la gamma. La risposta di KVQuant è l’ordinamento: quantizzare le chiavi prima che la rotazione venga applicata e applicare RoPE dopo la dequantizzazione. Accanto alla quantizzazione delle chiavi per canale, ai tipi di dati non uniformi e all’isolamento di una piccola frazione di outlier, questo li porta sotto il 0,1 di degrado della perplessità a 3 bit e consente di servire LLaMA-7B fino a 1 milione di token di contesto su un singolo A100-80GB.

È anche importante capire il livello di impatto da RoPE. Gli autori del paper “RotateKV” hanno segnalato un aumento del 145% negli errori di quantizzazione una volta aggiunto RoPE e hanno notato che i canali outlier differiscono tra teste di attenzione – il che è il motivo per cui l’applicazione di una matrice di rotazione condivisa ovunque è insufficiente e le rotazioni adattive ai capi funzionano meglio.

La tassa dei sistemi e perché non è un dettaglio

La quantizzazione per token si adatta bene alla decodifica. Ogni token arriva; lo si quantizza, lo si aggiunge alla sequenza (lungo la dimensione del token), nulla altro si muove.

Tuttavia, la quantizzazione per canale non si adatta. Poiché le statistiche di un canale si estendono ai token che non sono ancora stati generati, non si può calcolare un fattore di scala quando arriva un token. Il workaround di KIVI è quello di mantenere i token più recenti – fino a 128 – in precisione completa in un buffer residuo e quantizzare in gruppi una volta che se ne sono accumulati abbastanza.

Come si verifica, il buffer residuo diventa portante, piuttosto che solo una cosa incidentale. Su GSM8K con Llama-2-7B, i punteggi di precisione completa sono 13,50. Pienamente quantizzato a 2 bit con gli assi corretti, punteggia 5,76. Gli stessi assi e gli stessi bit, più il buffer residuo dei token recentemente prodotti in precisione completa, punteggiano 12,74. Una finestra mobile di token recentemente prodotti in precisione completa può recuperare gran parte di ciò che è stato perso a causa della quantizzazione aggressiva su problemi multi-step difficili – il che avrebbe senso se consideriamo quali token stanno venendo attenzionati da una catena di operazioni aritmetiche.

C’è un beneficio significativo nel fare tutte queste cose correttamente – come segnalato da KIVI, 2,6 volte meno utilizzo di memoria di picco per Llama-2-7B, consentendo dimensioni del batch fino a 4 volte più grandi, nonché un throughput 2,35-3,47 volte migliore in un compito di servizio reale.

Cosa fare con questo

1. Non utilizzare un solo quantizzatore per entrambi. Utilizzare quantizzatori diversi per le chiavi (per canale) e per i valori (per token). Una pipeline che applica un solo quantizzatore alla “cache KV” ha probabilmente già sacrificato la maggior parte della qualità possibile quando si utilizza un piccolo numero di bit per rappresentare ogni valore.
2. Quantizzare le chiavi prima di RoPE. Questo è un problema di correttezza piuttosto che una questione di preferenza.
3. Memorizzare una finestra di precisione completa di token recentemente generati. Sebbene memorizzare una tale finestra richieda poca memoria rispetto a quanto possa essere grande una cache, è proprio quest’area che genera gran parte dell’accuratezza per i compiti difficili.
4. Non convalidare sull’errore di ricostruzione. Convalidare sempre in base all’output dell’attenzione o alle prestazioni del compito finale. La metrica di archiviazione non è solo rumorosa – per i valori punta nella direzione sbagliata.
5. Non convalidare sui benchmark di scelta multipla a contesto breve. Gli autori di KIVI evitano intenzionalmente compiti chiusi come MMLU per questa valutazione, poiché un singolo passo di decodifica che legge i logaritmi di output a malapena esercita la cache. Qualsiasi valutazione che non costruisca una cache nel tempo e non esegua la generazione da essa non sarà in grado di osservare i fallimenti inerenti alla progettazione del sistema.

Dove è diretto il lavoro

Sebbene ci sia ancora qualcosa da fare riguardo alla natura geometrica del problema, molti ricercatori continuano a studiare modi in cui gli outlier di canali sono distribuiti tra le diverse teste del trasformatore e come le limitazioni hardware influenzano quali raggruppamenti sono più economici: InnerQ piega la normalizzazione dei canali delle chiavi all’interno dei pesi delle chiavi e delle query durante la prefilla. Pertanto, non si verifica alcun sovraccarico aggiuntivo in fase di esecuzione. Inoltre, InnerQ memorizza finestre di precisione elevate per i token recentemente generati e i token dell’attenzione. Facendo ciò, InnerQ elimina l’opportunità per gli outlier nel canale di contaminare i canali vicini.

Altri propongono che invece di memorizzare l’intera cache, dovremmo memorizzare solo abbastanza informazioni per poter rimaterializzare la chiave e/o il valore su richiesta da una rappresentazione della cache più piccola.

Infine, è importante ricordare che l’accuratezza non è l’unico parametro che la quantizzazione influenza. La ricerca pubblicata di recente ha dimostrato un degrado dell’allineamento a causa della quantizzazione delle cache KV. Inoltre, questa ricerca ha documentato un degrado dell’allineamento anche in ambienti di servizio vLLM di produzione che utilizzano cache FP8 insieme a un protocollo di recupero senza formazione che ha ripristinato fino al 97% di ciò che è stato perso in termini di allineamento. Pertanto, mentre una configurazione può mantenere i propri risultati di benchmark, non significa necessariamente che mantenga tutti gli altri parametri rilevanti di cui ci si preoccupa.

Il principio generale

L’idea di quantizzazione è stata inquadrata come un “budget di precisione”: quanti bit posso permettermi di sacrificare? La cache KV mostra che la domanda più utile è strutturale. La precisione è allocata in gruppi; il gruppo è l’unità di danno e l’asse che si raggruppa determina quali elementi condividono il loro destino. L’asse corretto è quello su cui il tensore viene consumato, cioè il modo in cui si utilizza il tensore e non come il tensore appare quando viene archiviato in memoria. Le chiavi vengono utilizzate tramite un calcolo del prodotto scalare contro la query. Un canale corrotto singolo avvelenerà tutti i punteggi. I valori vengono consumati attraverso un calcolo di media ponderata sparsa attraverso i token. Pertanto, un token corrotto singolo viene semplicemente pesato. Due tensori di dimensioni identiche e generati da due livelli consecutivi vengono trattati in modo diverso. Vale la pena chiedersi per ogni attivazione che si desidera comprimere: quale operazione contrae questo e il mio raggruppamento lo rispetta?

Due tensori di dimensioni identiche e generati da due livelli consecutivi vengono trattati in modo diverso. Vale la pena chiedersi per ogni attivazione che si desidera comprimere: quale operazione contrae questo e il mio raggruppamento lo rispetta? Quindi, un token corrotto singolo è semplicemente pesato. Due tensori di dimensioni identiche e generati da due livelli consecutivi vengono trattati in modo diverso. È importante chiedersi per ogni attivazione che si desidera comprimere: quale operazione contrae questo e il mio raggruppamento lo rispetta? Quindi, un token corrotto singolo è semplicemente pesato. Due tensori di dimensioni identiche e generati da due livelli consecutivi vengono trattati in modo diverso. Vale la pena chiedersi per ogni attivazione che si desidera comprimere: quale operazione contrae questo e il mio raggruppamento lo rispetta? Quindi, un token corrotto singolo è semplicemente pesato. Due tensori di dimensioni identiche e generati da due livelli consecutivi vengono trattati in modo diverso. Vale la pena chiedersi per ogni attivazione che si desidera comprimere: quale operazione contrae questo e il mio raggruppamento lo rispetta? Quindi, un token corrotto singolo è semplicemente pesato.