The Pulse
AWS serve Laplace 3B di i14 con uno stato fisso di 29,5 MB
AWS afferma che, su Trainium2, lo stato ricorrente usato per l’inferenza dal modello Laplace 3B di i14 resta di 29,5 MB per sequenza, indipendentemente dalla lunghezza del contesto. L’azienda dichiara 11.544 token di output al secondo su 1.
AI.info Team · Scritto con l’IA · verificato ·
L’attenzione lineare comprime la cronologia anziché conservare ogni chiave e ogni valore.
Kevin Shabahang, Iman Abbasnejad, Massimo Bardetti, Daniel Wirjo e Xiaodong Yu, AWS Machine Learning Blog
Lo stato di Laplace 3B resta fisso durante l’inferenza
AWS e la società di ricerca sull’IA i14 affermano di aver addestrato e servito Laplace 3B, un modello ad attenzione lineare, su un’istanza AWS Trainium2, mantenendo lo stato ricorrente usato per l’inferenza a 29,5 megabyte per sequenza. A differenza della cache chiave-valore di un transformer, che cresce man mano che si aggiunge contesto a una conversazione, lo stato di Laplace mantiene le stesse dimensioni anche quando aumenta la lunghezza del contesto.
Le aziende dichiarano 11.544 token di output al secondo su 1.280 sequenze simultanee in una singola istanza trn2.48xlarge. AWS afferma che il modello è stato eseguito tramite un motore vLLM non modificato, con un wrapper che collega il suo stato ricorrente al pool di memoria di vLLM. Il risultato è una dimostrazione tecnica, non un confronto pubblicato con un modello transformer sulla qualità delle risposte o sui costi di inferenza.
GLA sostituisce i token conservati con uno stato ricorrente
Laplace 3B ha 24 livelli decoder e 3,028 miliardi di parametri. Ogni livello sostituisce l’attenzione softmax con Simple Gated Linear Attention, che aggiorna uno stato matriciale di dimensioni fisse mentre elabora ciascun token. AWS descrive il calcolo dell’attenzione durante l’addestramento come un processo che cresce linearmente con la lunghezza della sequenza, anziché quadraticamente, e lo stato usato per l’inferenza come indipendente dalla lunghezza del contesto.
Questa compressione cambia ciò che il modello conserva: riassume il contesto precedente invece di mantenere ogni chiave e ogni valore. AWS avverte che i team dovrebbero valutare la qualità del recupero delle informazioni per i propri carichi di lavoro. Il post indica come possibili impieghi gli agenti attivi per lunghi periodi e l’analisi di documenti di grandi dimensioni, ma non riporta una valutazione per attività specifiche che mostri le prestazioni di Laplace su questi carichi di lavoro.
Trainium2 ha richiesto modifiche ai kernel e alla comunicazione
L’implementazione originale di Laplace utilizzava kernel Triton ottimizzati per le GPU. Il team dell’AWS Generative AI Innovation Center ha sostituito il kernel di Gated Linear Attention con operazioni PyTorch standard, organizzate in 16 blocchi da 128 token, e ha usato il backend TorchNeuron Native in versione beta per Trainium. La dimensione dei blocchi corrisponde alle 128 partizioni del buffer di stato integrato in ciascun NeuronCore.
L’addestramento ha richiesto anche di modificare il modo in cui l’ottimizzatore Muon comunicava tra i 64 NeuronCore logici. AWS afferma che due operazioni all-to-all hanno ridotto di 24 volte il traffico di dati per rank e aumentato il throughput di 1,63 volte, da 73.378 a 119.934 token al secondo. Il team ha compilato il passaggio all’indietro sequenziale svolgendo i suoi 16 passaggi ricorrenti per rientrare nei limiti del compilatore relativi al grafo e alla memoria.
I risultati dichiarati hanno limiti precisi
Il piano di addestramento punta a 100 miliardi di token FineWeb-Edu, ma i risultati riportati nel post coprono 84.210 passaggi, pari a circa 44 miliardi di token. AWS dichiara un throughput di circa 120.000 token al secondo durante quell’esecuzione. Lo stato usato per l’inferenza, secondo i dati riportati, è circa 82 volte più piccolo di una cache KV comparabile con 8.192 token; nel test del team, portare il prompt da 2.048 a 8.192 token ha aumentato dello 0,5% il costo per token di output.
Questi dati riguardano un modello, una singola istanza di acceleratore e misurazioni effettuate da AWS; il post non fornisce un confronto diretto con un transformer usato come riferimento né un confronto dei prezzi. Individua inoltre una complicazione pratica nell’inferenza: quando cambia la larghezza dei batch in ingresso, le nuove forme possono richiedere una compilazione che, secondo il team, richiede circa 11 secondi al primo utilizzo. Per Laplace 3B, la memoria fissa elimina l’aumento dovuto alla lunghezza del contesto dal calcolo dello stato per sequenza, mentre il comportamento nel recupero delle informazioni e il traffico variabile richiedono ancora test specifici per ciascun carico di lavoro.