Riepilogo Tecnico: Pre-addestramento di Modelli Linguistici di Grandi Dimensioni con MXFP4

Enunciato del Problema

La rapida scalabilità dei Modelli Linguistici di Grandi Dimensioni (LLM) ha intensificato la domanda di algoritmi di addestramento efficienti dal punto di vista hardware. Sebbene FP8 si sia affermato come un'alternativa a bassa precisione praticabile, l'addestramento con formati a 4 bit (come MXFP4) rimane una sfida significativa. Una difficoltà primaria è che gli outlier nelle attivazioni e nei gradienti amplificano gli errori di quantizzazione, potenzialmente destabilizzando l'ottimizzazione. Studi preliminari indicano che la quantizzazione grossolana degli outlier nelle attivazioni può dominare l'intervallo di quantizzazione, degradando severamente la qualità in virgola mobile. Nonostante gli sforzi per mitigare il rumore tramite precisione mista o schedule a stadi, rimane una domanda fondamentale: Perché l'addestramento FP4 a pipeline completa diverge spesso anche quando le attivazioni forward e i gradienti delle attivazioni appaiono stabili? Inoltre, non è chiaro quali componenti specifici della pipeline di addestramento (propagazione forward, gradienti delle attivazioni o gradienti dei pesi) siano più responsabili del degrado della convergenza e quali meccanismi di stabilizzazione siano efficaci.

Metodologia

Gli autori hanno condotto uno studio controllato sul pre-addestramento di Llama 3.1–8B utilizzando il dataset C4, aderendo alla configurazione di pre-addestramento LLM di MLPerf. Lo studio è stato eseguito su GPU AMD Instinct MI355X sfruttando il supporto nativo per tensor-core FP4, distinguendo questo lavoro dalle ricerche precedenti che si basavano su emulazione software.

Il disegno sperimentale centrale ha coinvolto una "scala" progressiva di abilitazione di MXFP4:

1. Baseline: Addestramento completo in FP8.
2. Quantizzazione a stadi: MXFP4 è stato abilitato progressivamente su tre componenti, mantenendo gli altri fissi in FP8:
   • Solo propagazione forward (Fprop).
   • Fprop + gradienti delle attivazioni (Dgrad).
   • Fprop + Dgrad + gradienti dei pesi (Wgrad).
3. Valutazione della stabilizzazione: Sotto la pipeline MXFP4 completa (Fprop + Dgrad + Wgrad), gli autori hanno testato varie interventi per ripristinare la stabilità:
   • Arrotondamento Stocastico: Per introdurre casualità.
   • Rotazioni Hadamard Randomizzate: Per decorrelare gli outlier.
   • Rotazioni Hadamard Deterministiche: Utilizzando specificamente matrici $16 \times 16$ (H16) e $32 \times 32$ (H32).

Metriche: La convergenza è stata definita come il raggiungimento di una perplessità di validazione di 3.3. La metrica primaria per l'efficienza è stata il sovra-costo in token (l'aumento relativo dei token richiesti per raggiungere l'obiettivo rispetto alla baseline FP8). Le metriche secondarie includevano il throughput degli step di addestramento e l'accelerazione end-to-end.

Contributi e Risultati Chiave

1. Identificazione della Fonte di Instabilità

Lo studio isola la quantizzazione del Gradiente dei Pesi (Wgrad) come il principale motore del degrado della convergenza.

• Abilitare MXFP4 solo per Fprop ha comportato un sovra-costo in token dell'8–9%.
• Abilitare Fprop + Dgrad ha prodotto un sovra-costo simile del 10–11%.
• Abilitare la pipeline completa (Fprop + Dgrad + Wgrad) senza stabilizzazione ha fatto saltare il sovra-costo in token al 26–27% o ha portato a una divergenza completa ("Non converge").
• Conclusione: L'instabilità non è intrinseca alla quantizzazione di Fprop o Dgrad, ma è dominata dalla quantizzazione di Wgrad.

2. Valutazione delle Strategie di Stabilizzazione

Gli autori hanno testato se l'aggiunta di stocasticità o rotazioni strutturate potesse stabilizzare la pipeline completa:

• Arrotondamento Stocastico: Non è riuscito a stabilizzare la pipeline completa; l'esecuzione ha diverso.
• Rotazioni Hadamard Randomizzate: Analogamente non sono riuscite a convergere quando applicate alla pipeline completa.
• Rotazioni Hadamard Deterministiche: Hanno ripristinato con successo un'ottimizzazione stabile. Applicare una trasformazione Hadamard deterministica (specificamente H16) alla pipeline MXFP4 completa ha ridotto il sovra-costo in token nuovamente all'8–9%, tracciando da vicino la baseline FP8.
• Insight: Gli interventi stocastici sembrano amplificare l'errore di quantizzazione effettivo piuttosto che mitigare gli effetti degli outlier. L'instabilità è guidata da errori di micro-scaling strutturati lungo percorsi di gradiente sensibili, che le rotazioni deterministiche possono correggere.

3. Efficienza e Throughput

Con il problema di stabilità risolto tramite Hadamard Deterministico (H16), gli autori hanno dimostrato i guadagni pratici di efficienza di FP4:

• Throughput: La configurazione MXFP4 + H16 ha raggiunto un throughput degli step di addestramento superiore del 20% rispetto alla baseline FP8 sulle GPU MI355X.
• Accelerazione End-to-End: Poiché il sovra-costo in token è stato minimizzato (8–9%), il vantaggio di throughput si è tradotto in un'accelerazione end-to-end del 9–10%.
• Insight Chiave: Le accelerazioni FP4 sono "governate dalla stabilità". Senza controllare l'instabilità di Wgrad, l'aumentato requisito di token per convergere nega i benefici computazionali grezzi della bassa precisione.

Significato e Affermazioni

Il documento afferma di fornire la prima dimostrazione di pre-addestramento completo MXFP4 su hardware di classe MI355X senza emulazione software. Il suo significato risiede nel cambiare la comprensione della dinamica dell'addestramento FP4:

• Meccanismo di Fallimento: L'instabilità nel pre-addestramento FP4 non è dovuta a una stocasticità insufficiente, ma piuttosto a errori di micro-scaling strutturati in percorsi di gradiente sensibili (specificamente Wgrad).
• Ricetta Pratica: Una rotazione Hadamard deterministica è uno stabilizzatore necessario ed efficace per l'addestramento MXFP4 a pipeline completa, consentendo la realizzazione di guadagni di efficienza hardware.
• Limitazione Contestuale: Gli autori notano modestamente che questa ricetta di stabilizzazione è dipendente dal contesto. Sebbene efficace per Llama 3.1–8B sul dataset C4, il comportamento può variare tra modelli, dataset o metodi di fine-tuning diversi, e non dovrebbe essere assunto universale senza ulteriori evidenze.

In sintesi, il lavoro stabilisce che la sostituzione di FP8 con MXFP4 è fattibile per il pre-addestramento su larga scala, a condizione che gli errori di quantizzazione di Wgrad siano controllati tramite interventi strutturali deterministici piuttosto che tramite iniezione di rumore stocastico.