Un nuovo trucco riduce drasticamente i costi della compressione dei modelli AI giganti

Punti chiave

• La distillazione della conoscenza, che copia il comportamento di un grande modello AI in uno più piccolo, tipicamente richiede centinaia di chip specializzati e un coordinamento attento.
• Un nuovo articolo memorizza in cache gli output del grande modello una sola volta all'inizio, quindi non deve mai essere eseguito di nuovo durante l'addestramento, riducendo significativamente l'uso della memoria.
• Una seconda tecnica, chiamata fused chunked KL loss, elabora i dati di addestramento in piccole sezioni invece che tutto in una volta, mantenendo la memoria di picco quasi stabile anche con input molto lunghi.
• A 32.000 token di contesto, la memoria di picco scende da 85,2 gigabyte a 5,45 gigabyte, una riduzione di 15,6 volte.
• Il codice completo per la funzione di loss efficiente in termini di memoria è disponibile pubblicamente su GitHub.

I modelli AI più grandi al mondo sono quasi impossibilmente voluminosi. Kimi-K3, rilasciato all'inizio di quest'anno, ha 2,8 trilioni di parametri (pensa ai parametri come i milioni di piccoli dial all'interno di un modello che vengono regolati durante l'addestramento) e necessita di circa tre terabyte di memoria specializzata solo per il caricamento. Questo equivale a uno spazio superiore a quello di venti schede grafiche consumer di fascia alta.

Poiché l'esecuzione di modelli di queste dimensioni è così costosa, i ricercatori li comprimono regolarmente in versioni più piccole che costano meno da gestire. Il metodo standard è chiamato distillazione della conoscenza: esegui il grande modello "insegnante" su un mucchio di testo, registra come risponderebbe, quindi addestra un modello "studente" più piccolo per copiare quelle risposte il più fedelmente possibile. Lo studente finisce per essere molto più economico da eseguire mantenendo la maggior parte della qualità dell'insegnante.

Perché è così difficile farlo a basso costo?

La spesa proviene da due fonti. Innanzitutto, l'approccio standard mantiene sia l'insegnante che lo studente caricati in memoria contemporaneamente. In secondo luogo, il confronto dei loro output richiede la costruzione di una griglia enorme: una riga per ogni parola nel vocabolario del modello (spesso più di 100.000 parole), una colonna per ogni posizione nell'input. Moltiplicati insieme in un lungo pezzo di testo e riempi rapidamente ogni chip che hai.

Come esempio concreto: un modello open popolare, gpt-oss-120b, ha un vocabolario di 201.088 parole. Con una lunghezza di sequenza di 32.000 token (all'incirca il testo di un breve romanzo) e una modesta dimensione di batch di quattro esempi alla volta, il tensore di output dell'insegnante da solo occupa circa 50 gigabyte. Aggiungi il modello studente, i gradienti e tutto il resto di cui il processo di addestramento ha bisogno, e un singolo passo di addestramento può raggiungere circa 250 gigabyte, più di quanto un H200, uno dei chip AI più potenti disponibili oggi, possa contenere.

Cosa hanno cambiato i ricercatori?

Due cose, che lavorano insieme. Entrambe sono descritte in un articolo pubblicato dal team di CompactifAI, condiviso per la prima volta su Hugging Face.

Il primo cambiamento è la distillazione offline. Invece di rieseguire l'insegnante ad ogni passo di addestramento, i ricercatori lo eseguono una volta e salvano solo le sue 100 previsioni di parole più probabili per posizione. Quella cache viene scritta su disco. L'insegnante viene quindi disattivato completamente e lo studente si addestra rispetto alla cache salvata da solo. Poiché la cache può essere riutilizzata in molti esperimenti diversi, l'esecuzione costosa dell'insegnante diventa un costo una tantum.

Il secondo cambiamento è la fused chunked KL loss. La divergenza KL è una misura matematica di quanto due distribuzioni di probabilità siano diverse, qui, di quanto le previsioni di parole dello studente siano diverse da quelle dell'insegnante. Il modo standard di calcolarla costruisce quella griglia vocabolario-per-sequenza completa in una volta e la mantiene tutta in memoria. Il nuovo metodo divide la sequenza in piccoli chunk, calcola e scarta ogni chunk prima di passare al successivo, e ricalcola ciò di cui ha bisogno durante il pass all'indietro piuttosto che immagazzinarlo.

Metodo	Memoria di picco (contesto 8K)	Tempo di iterazione
Distillazione online	102,8 GB	25,9 s
Offline, KL denso	78,3 GB	18,5 s
Offline, KL forward-chunked	61,8 GB	18,4 s
Offline, fused chunked KL	58,3 GB	20,2 s

Tutti e quattro i metodi producono risultati di addestramento quasi identici, il che significa che la perdita di qualità derivante dall'utilizzo delle previsioni top-100 memorizzate in cache invece della distribuzione completa dell'insegnante è trascurabile.

Questo è importante per documenti più lunghi?

Sì, e il divario cresce rapidamente. Il vantaggio di memoria dell'approccio fused chunked è modesto a 8.000 token ma drammatico a lunghezze più lunghe. A 32.000 token, la memoria di picco scende da 85,2 gigabyte con l'approccio standard a 5,45 gigabyte con quello nuovo. La loss densa fallisce completamente a 64.000 token. A 256.000 token, il metodo fused chunked utilizza 11,6 gigabyte rispetto a 134,2 gigabyte per l'alternativa più vicina.

In termini semplici: gli esperimenti che in precedenza richiedevano un rack di centinaia di chip possono ora essere eseguiti su una singola scheda di fascia alta. Ciò rende realistico per i team di ricerca più piccoli e le aziende costruire modelli compressi dai più grandi sistemi AI open-source disponibili oggi.

Domande comuni

Il modello più piccolo finisce per essere peggiore dell'originale?

Nei benchmark riportati, la perdita di addestramento del modello studente è stata quasi identica in tutti e quattro i metodi, incluso il nuovo approccio economico. La qualità si è mantenuta nonostante l'utilizzo di sole 100 previsioni di parole per posizione invece dell'output completo dell'insegnante.

Chiunque può usare questa tecnica?

Il codice della chunked loss è open source e pubblicato pubblicamente su GitHub da CompactifAI. Chiunque abbia un singolo chip AI moderno e una copia di un grande modello open-source può, in linea di principio, eseguire i propri esperimenti di distillazione con esso.

Chi beneficia per primo?

I piccoli team di ricerca AI e le aziende che desiderano comprimere grandi modelli open-source per un deployment più economico avvertiranno la differenza più rapidamente. Gli utenti finali potrebbero beneficiare successivamente man mano che modelli compressi più accessibili raggiungono prodotti e servizi.