Articolo
Accelerare Sonar attraverso la speculazione
La decodifica speculativa accelera la velocità di generazione dei Modelli di Linguaggio di Grandi Dimensioni (LLM) utilizzando un modello di bozza rapido e ridotto per produrre candidati al completamento che vengono verificati dal più grande modello di destinazione.
Con questo schema, invece di una sola esecuzione del costoso modello di destinazione per produrre un singolo token, ne vengono emessi molteplici in un solo passaggio. Qui presentiamo i dettagli dell'implementazione di vari tipi di decodifica speculativa, applicata in Perplexity per ridurre la latenza tra token sui modelli Sonar.
Decodifica Speculativa
La Decodifica Speculativa sfrutta la struttura delle lingue naturali e la natura autoregressiva dei trasformatori per accelerare la generazione di token. Anche se modelli più grandi, come Llama-70B, portano più conoscenza rispetto a quelli più piccoli, come Llama-1B, su alcuni compiti più semplici si comportano in modo simile. Questa sovrapposizione suggerisce che certe sequenze sono meglio generate da modelli meno costosi, lasciando problemi complessi ai modelli più grandi. La sfida sta nel determinare quali completamenti sono migliori e se la generazione del modello più piccolo è di qualità pari a quella del modello più grande.
Fortunatamente, gli LLM sono trasformatori autoregressivi: quando ricevono una sequenza di token, producono la distribuzione di probabilità del token successivo. Inoltre, i logits derivati dalle caratteristiche intermedie associate ai token nella sequenza di input indicano anche quanto è probabile che il modello emetta proprio quei token. Questa proprietà consente la speculazione: se una sequenza di token è generata da uno più piccolo a partire da un prefisso di input, può essere eseguita attraverso quello più grande per determinare quanto bene si allinea con il modello di destinazione. Ogni prefisso dei candidati viene valutato con una probabilità e viene scelto quello più lungo sopra una soglia di accettazione. Come bonus, il modello di destinazione fornisce anche un token successivo gratuitamente: se un modello di bozza genera n token, fino a n + 1 possono essere emessi in un solo passaggio.
Al momento dell'inferenza, il processo di campionamento speculativo può essere suddiviso in circa 4 fasi:
Riempimento iniziale: sia i modelli di destinazione che di bozza devono essere eseguiti sulla sequenza di input per popolare le voci della cache KV. Mentre alcuni schemi, come Medusa, utilizzano strati densi più semplici per la previsione, in questo post ci concentriamo sulle bozze basate su trasformatori che necessitano delle proprie cache KV.
Generazione di bozze: il modello di bozza itera per produrre un numero fisso di token. La sequenza di bozza può essere lineare o il modello può esplorare una struttura ad albero fino a una certa profondità (EAGLE, Medusa). Qui ci concentriamo su sequenze lineari.
Accettazione: il modello di destinazione funziona sulla sequenza di bozza, costruendo logits corrispondenti a ciascun token di bozza. Si determina la lunghezza della sequenza più lunga accettabile.
Generazione di destinazione: poiché i logits generati dal modello di destinazione, nella posizione non corrispondente o alla fine della sequenza corrispondono a un token successivo. Questi logits possono essere campionati per fornire un token robusto dalla destinazione, completando la sequenza.
Esistono vari metodi per implementare la decodifica speculativa. In questo post, ci concentreremo sui metodi che abbiamo usato per accelerare i modelli Sonar utilizzando un modello interno da 1 B, così come sui meccanismi di previsione che stiamo costruendo per accelerare modelli alla scala di DeepSeek.
Target-Draft
La decodifica speculativa si può ottenere accoppiando un piccolo LLM esistente come modello di bozza a un modello di destinazione per generare sequenze di candidati. In produzione, abbiamo accelerato Sonar usando un modello Llama-1B ottimizzato sullo stesso dataset del target. Sebbene questo approccio non richiedesse l'addestramento di una bozza da zero, il modello piccolo utilizza comunque una significativa capacità di cache KV e introduce un leggero overhead di riempimento, aumentando il TTFT.
Sotto questo schema, il decodificatore specula solo su batch di sola decodifica, generando token tramite campionamento standard durante il riempimento o su batch misti di riempimento-decodifica. Nella fase di riempimento, i logits di destinazione vengono immediatamente campionati per precompilare anche il token appena generato nella cache KV della bozza. La bozza non viene ancora campionata, ma i logits che produce vengono trasferiti alla fase di decodifica.
In fase di decodifica, il modello di bozza avanza, campionando il token principale in ogni fase. Una volta raggiunta la lunghezza desiderata della bozza, i token vengono eseguiti tramite il modello di destinazione per produrre i logits grazie ai quali il campionatore identifica la lunghezza della sequenza accettata. L'accettazione viene determinata confrontando le distribuzioni di probabilità complete tra la bozza e la destinazione. Poiché la destinazione produce sempre un set di logits dopo la sequenza di bozza accettata, questo è campionato per produrre un output aggiuntivo. Poiché il modello di bozza non ha ancora visto quel token accettato, viene rieseguito per popolare le corrispondenti voci della cache KV in preparazione per il prossimo passo di decodifica, riportando nuovamente i logits.
EAGLE
EAGLE è uno schema di decodifica speculativa che esplora più sequenze di bozza, generate attraverso una traversal di albero di token di bozza probabili. Un albero fisso (EAGLE) o a forma dinamica (EAGLE-2) viene esplorato utilizzando esecuzioni consecutive dei token di bozza, considerando i Top-K candidati in ciascun nodo anziché seguire il token con punteggio più alto in una sequenza lineare. Le sequenze vengono quindi valutate e viene selezionata la più lunga adatta per continuare, aggiungendo anche un token aggiuntivo derivato dal target.
Per ottenere una previsione più precisa, un modello di bozza EAGLE prevede non solo in base ai token, ma utilizzando anche le caratteristiche target (stati nascosti dell'ultimo strato) del modello di destinazione. Lo svantaggio di EAGLE è la necessità di allenare modelli di bozza personalizzati e ridotti, che siano sufficientemente precisi da generare candidati adatti entro un budget di bassa latenza. Tipicamente, un modello di bozza è un singolo strato trasformativo identico a uno strato di decodifica del modello originale, strettamente accoppiato al target attraverso l'associazione alle sue proiezioni di embedding e lm_head. Poiché ciò richiede meno capacità di cache KV, EAGLE ha un ingombro di memoria inferiore.
Per verificare sequenze ad albero nel modello di destinazione, devono essere usate maschere di attenzione personalizzate. Sfortunatamente, l'uso di una maschera di attenzione personalizzata per un'intera sequenza rallenta notevolmente l'attenzione per lunghezze di input realistiche (fino al 50%), annullando parte dell'accelerazione ottenibile attraverso la speculazione. Non abbiamo ancora implementato appieno l'esplorazione ad albero in produzione per questo motivo, concentrandoci invece sul caso speciale della previsione a singolo token tramite schemi simili a MTP presentati nel Rapporto Tecnico DeepSeek-V3.
MTP
Questo schema è simile alla decodifica bozza-destinazione, con l'eccezione che gli stati nascosti sono usati insieme ai token per la previsione. È necessario fare un po' più di lavoro sia nella fase di riempimento che di decodifica rispetto alla speculazione bozza-destinazione regolare. Il modello di bozza utilizza sia i token che gli stati nascosti: il token t_{i+1} viene campionato dai logits L_i corrispondenti al token t_i, derivati dagli stati nascosti H_i. Di conseguenza, i buffer dei token di input devono essere spostati di uno step a sinistra rispetto ai vettori di stato nascosto emessi dal target. La figura sotto indica le corrispondenze usate per l'addestramento, oltre allo spostamento durante l'inferenza.
Il flusso di decodifica è abbastanza simile alla decodifica bozza-destinazione, con l'eccezione che sia gli stati nascosti che i logits vengono trasferiti. La nostra implementazione condivide tutti i sampling associati e l'elaborazione dei logits, specializzando solo le invocazioni di avanzamento del modello. Quando vengono previsti più token, il modello di bozza utilizza gli stati nascosti di bozza per la previsione, popolando anche le voci della cache KV basandosi sulle proprie caratteristiche. A lungo termine, questo può degradare la precisione. Successivamente, quando si esegue il modello di bozza per popolare la voce della cache KV per la previsione target, lo eseguiamo sull'intera sequenza prendendo gli stati target nascosti più accurati come input. Poiché questi modelli di bozza sono piccoli, il costo aggiuntivo di elaborare i token aggiuntivi è trascurabile.
Training delle Teste MTP
Per beneficiare dell'MTP, abbiamo costruito l'infrastruttura necessaria per addestrare teste MTP attaccate ai nostri modelli finemente sintonizzati sui dataset di Perplexity, funzionando su un nodo con 8 dispositivi H100. In circa un giorno, possiamo costruire teste per modelli che vanno da Llama-1B a Llama-70B e DeepSeek V2-Lite. Per modelli più grandi, ci affidiamo a teste MTP costruite durante il processo di perfezionamento.
L'obiettivo dell'addestramento MTP è abbinare gli stati nascosti della bozza e i logits estrapolati dagli stati nascosti target ai logits e stati nascosti del token successivo del target. Poiché l'inferenza per stati nascosti è costosa, li pre-calcoliamo utilizzando la nostra implementazione di inferenza ottimizzata del modello target, per essere usati durante l'addestramento. Tuttavia, per convalidare l'implementazione di inferenza MTP e garantire che differenze numeriche dovute a quantizzazione o ottimizzazioni non ostacolino i risultati, per la stima della perdita di validazione e dell'accuratezza riutilizziamo completamente l'implementazione di inferenza di entrambi i modelli target e bozza.
Quando si scala dal dataset ShareGPT utilizzato nell'articolo originale a campioni più grandi, abbiamo notato che l'architettura della testa MTP delineata e implementata nel documento EAGLE non riusciva a addestrare per modelli delle dimensioni di 70B. A differenza di ShareGPT che conteneva un numero maggiore di sequenze più corte, alleniamo su un numero leggermente inferiore di prompt sostanzialmente più lunghi. Poiché le teste EAGLE originali differivano leggermente nella struttura da un tipico trasformatore, abbiamo reintrodotto alcuni strati di normalizzazione RMS che erano stati rimossi. Abbiamo trovato che ciò non solo permetteva all'addestramento di convergere, ma migliorava anche l'accuratezza delle teste di qualche punto percentuale.
Non solo i normali strati facilitano l'addestramento, la reintroduzione delle norme è anche matematicamente intuitiva. Le teste MTP riutilizzano gli embedding e le proiezioni dei logits del modello target, poiché possono essere sostanziali in dimensione (circa 2 GB per Llama 70B). Durante l'addestramento, questi sono congelati e l'aspettativa è che lo strato MTP impari a inglobare previsioni nello stesso spazio vettoriale di ciò che lo strato di proiezione del modello originale ha appreso durante l'addestramento. Eliminando le norme, ci si aspetta che un singolo MLP apprenda la stessa funzione di un MLP seguito da una norma, che ostacola l'abbinamento tra gli stati nascosti dei modelli di bozza e di target.
Inferenza con Decodifica Speculativa
Nell'engine di inferenza, per generare token per sequenze di input, devono essere prima raggruppati in batch di dimensioni ragionevoli, quindi le pagine devono essere allocate nella cache KV per i prossimi token. I token di input e le informazioni sulla pagina KV vengono poi confezionati in un buffer trasmesso a tutti i ranghi paralleli che eseguono il modello. Infine, i metadati vengono copiati nella memoria GPU e il modello viene eseguito per produrre i logits da cui viene campionato il token successivo.
A differenza di certe implementazioni che collegano liberamente un server di inferenza bozza e target tramite un wrapper che orchestra le richieste tra di loro, le nostre coppie bozza-target sono strettamente collegate e attraversano la generazione all'unisono. La pianificazione dei batch e l'allocazione delle pagine KV è condivisa tra i modelli per tutte le forme di decodifica speculativa: questo unifica la logica che connette un modello al server di inferenza complessivo, poiché tutti espongono la stessa interfaccia.
Il runtime di inferenza in Perplexity è modellato attorno a FlashInfer, che determina i metadati che devono essere costruiti per configurare e pianificare il kernel di attenzione. Date alcune sequenze di input che formano un batch, per il riempimento, la decodifica o la verifica, deve essere svolto del lavoro lato CPU per allocare buffer intermedi e popolari certi buffer costanti usati nell'attenzione. Questo lavoro si aggiunge al costo della pianificazione dei batch e all'allocazione delle pagine KV, che comportano anche latenze che devono essere nascoste per massimizzare l'utilizzo della GPU.
Mentre abbiamo completamente parallelizzato il lavoro lato CPU e GPU per l'inferenza senza speculazione, abbiamo scoperto che l'equilibrio CPU-GPU per la decodifica speculativa è più complesso. La sfida principale sorge dal fatto che il numero di token accettati determina la lunghezza della sequenza per una successiva esecuzione, introducendo un punto di sincronizzazione GPU-CPU difficile da evitare. Abbiamo sperimentato diverse pianificazioni per nascondere al meglio la latenza del lavoro della CPU.
Pianificazione Bozza-Target
Nonostante sia più piccolo di un modello target, quando un intero LLM è usato come bozza, introduce ancora una considerevole latenza sulla GPU, fornendo un certo margine per nascondere operazioni CPU costose. Poiché i modelli più piccoli non beneficiano del parallelismo tensore, c'è un disallineamento tra il numero di ranghi che un target e una bozza sono partizionati. Nella nostra implementazione, il modello di bozza viene eseguito solo sul nodo leader di un gruppo TP.
Come indicato in precedenza, un passo di decodifica trasferisce logits alla prossima esecuzione. Questo consente di sovrapporre un'esecuzione del modello di bozza con il lavoro di pianificazione dei batch lato CPU. Dopo che il batch viene messo insieme, chiamate ripetute al campionatore e alla bozza producono i token di bozza. In parallelo, il batch per la verifica viene messo insieme per il modello di destinazione e sincronizzato con i lavoratori paralleli. I logits di destinazione sono verificati e campionati per determinare le lunghezze delle sequenze accettate. A questo punto, è necessaria una sincronizzazione GPU-CPU per determinare le lunghezze delle sequenze successive. Poiché il modello di bozza è eseguito solo sul nodo leader, il batch è impostato in modo sequenziale e la sua esecuzione è avviata per popolare le voci della cache KV con il token aggiuntivo prodotto dal target. I logits prodotti da questa esecuzione di bozza nell'attuale esecuzione saranno usati per campionare il primo token di bozza nell'esecuzione successiva. Fondamentalmente, mentre la bozza è in esecuzione, il batch successivo può essere pianificato.
Pianificazione MTP per un Singolo Token
Mentre il runtime non fornisce ancora un'esplorazione in stile Eagle dei draft tree, abbiamo implementato un caso speciale di questo schema, considerando una sequenza lineare di token di bozza prodotti da un modello della dimensione di un singolo strato decodificatore del trasformatore. Questo schema può essere utilizzato per la previsione di bozza usando i pesi open-source di DeepSeek R1. Il sotto-caso della previsione di un singolo token è interessante, poiché strati MTP grandi ottengono tassi di accettazione sufficientemente alti da giustificare il loro overhead.
La pianificazione MTP è un po' più complessa, poiché il modello di bozza è molto più veloce, nascondendo meno latenza lato CPU. Inoltre, la bozza è partizionata insieme al modello di destinazione, richiedendo trasferimenti di memoria condivisa per le informazioni sui batch. Un'esecuzione inizia trasferendo le informazioni sui batch e campionando il primo token dai logits trasmessi, similmente allo schema precedente. Successivamente, il target viene eseguito per validare i token, elaborando 2 * D token, dove D è la dimensione del batch di decodifica. Questo è ideale per il micro-batching nei modelli Mixture-of-Experts (MoE) su interconnessioni più lente come InfiniBand, poiché il batch si divide equamente in due metà. Gli stati nascosti del target vengono trasferiti alla prossima esecuzione di bozza, mentre i logits sono passati al campionatore per la verifica.
Effettuando una quantità limitata di lavoro aggiuntivo sulla GPU, evitiamo la sincronizzazione CPU-GPU dopo l'accettazione della sequenza di bozza. Dopo che i token di input dei target vengono spostati, un kernel inserisce i prossimi token di destinazione nelle loro posizioni corrispondenti. La bozza viene quindi eseguita nuovamente con le stesse informazioni di batch del target, popolando le voci della cache KV e costruendo i logits e gli stati nascosti per la prossima esecuzione, facendo qualche lavoro ridondante su token che non sono stati accettati. In queste situazioni, la latenza del lavoro non utilizzato è appena misurabile a causa delle piccole dimensioni del modello di bozza. In parallelo all'esecuzione della bozza, le lunghezze delle sequenze vengono determinate sulla CPU e la pianificazione del batch successivo viene avviata, senza dover aspettare che il lavoro GPU termini.
L'overhead del lavoro aggiuntivo nello strato di bozza non è evidente nell'attenzione, tuttavia gli strati MLP sono più problematici. Poiché le istruzioni di moltiplicazione di matrici aggiustano a un limite di 64 lungo la dimensione del numero di token, se il raddoppio non richiede un numero significativamente maggiore di blocchi, l'overhead è nascosto. Per sequenze di bozza più lunghe l'overhead è più costoso e lo schema utilizzato per modelli bozza-destinazione regolari funziona meglio.
Riferimenti
Infezione Veloce dai Trasformatori tramite Decodifica Speculativa
EAGLE: Il Campionamento Speculativo Richiede di Riconsiderare l'Incertezza delle Caratteristiche
EAGLE-2: Inferenza Più Rapida dei Modelli Linguistici con Alberi di Bozza Dinamici
Medusa: Semplice Struttura di Accelerazione dell'Inferenza LLM con Teste Decodificatrici Multiple
FlashInfer: Motore di Attenzione Efficiente e Personalizzabile per il Servizio di Inferenza LLM