MXNorm: Reusing MXFP block scales for efficient tensor normalisation

Il paper MXNorm propone un'alternativa efficiente a RMSNorm che riutilizza le scale dei blocchi MXFP per ridurre drasticamente le operazioni di riduzione necessarie alla normalizzazione, ottenendo accelerazioni pratiche nei kernel e una precisione di addestramento quasi invariata per modelli come Llama 3.

L'essenziale

Immagina di dover preparare un enorme banchetto per un esercito di robot (le Intelligenze Artificiali) che devono imparare a parlare e ragionare.

Il Problema: Il Collo di Bottiglia della Cucina

Negli ultimi anni, i robot sono diventati incredibilmente veloci nel calcolare le moltiplicazioni (il "cuore" del loro apprendimento). È come se avessimo installato un motore da Formula 1 su un'auto: fanno un sacco di calcoli in un attimo.

Tuttavia, c'è un problema: mentre il motore è velocissimo, la cucina (la parte che prepara gli ingredienti) è rimasta lenta. Prima di poter cucinare, i robot devono "normalizzare" i dati. Immagina di dover pesare e ridistribuire equamente ogni singolo granello di sale in una zuppa gigante prima di poterla servire. Questo passaggio, chiamato RMSNorm, è lento e richiede molta energia, rallentando tutto il processo.

La Soluzione: MXNorm (Il Trucco del Cuoco)

Gli autori di questo studio (di Graphcore) hanno avuto un'idea geniale: perché fare due lavori separati quando puoi farne uno solo?

Attualmente, il processo è così:

1. Pesare la zuppa (Calcolare la media quadratica o RMS).
2. Ridistribuire il sale (Normalizzare).
3. Tagliare gli ingredienti in pezzi piccoli (Quantizzazione MXFP, per renderli più leggeri da trasportare).

MXNorm fa un trucco da mago: dice "Aspetta! Quando tagli gli ingredienti in pezzi piccoli (passo 3), ho già misurato quanto sono grandi i pezzi più grandi di quel blocco. Non devo pesare di nuovo tutta la zuppa!"

Invece di calcolare la media esatta di tutto, MXNorm riutilizza le misure già prese durante il taglio degli ingredienti per stimare quanto sale serve. È come se, invece di pesare ogni singolo granello di sale, guardassi il peso massimo di un cucchiaino e dicessi: "Ok, se questo è il massimo, la media sarà circa questa".

Come funziona in pratica?

1. Il Vecchio Metodo (RMSNorm): È come un ispettore di qualità molto preciso che conta ogni singolo granello di sale in una stanza enorme per assicurarsi che la distribuzione sia perfetta. È preciso, ma lento.
2. Il Nuovo Metodo (MXNorm): È un ispettore furbo che guarda solo i "pacchetti" di sale già pronti. Sa che se il pacchetto più grande pesa X, la media del pacchetto sarà circa Y. Non perde tempo a contare tutto, ma usa un'ottima stima basata sui pacchetti.

I Risultati: Più Veloce, Quasi Ugualmente Buono

Gli scienziati hanno provato questo metodo su modelli linguistici enormi (come Llama 3, che sono come enciclopedie viventi).

• Velocità: Hanno scoperto che MXNorm è fino a 2,4 volte più veloce in alcune operazioni specifiche. È come passare da un'auto a scatti a una che scorre fluida.
• Qualità: La "zuppa" (l'intelligenza del modello) è rimasta quasi identica. Per i modelli piccoli e medi, non c'è quasi nessuna differenza. Per i modelli giganti (8 miliardi di parametri), bisogna usare un tipo di stima leggermente diverso (chiamato "p=2", che è come usare una bilancia un po' più sensibile), ma il risultato finale è eccellente.

Perché è importante?

Oggi, i computer diventano sempre più potenti nel fare calcoli matematici complessi, ma la memoria e la velocità di trasferimento dei dati non tengono il passo.
MXNorm è come rimuovere un ingorgo da un'autostrada. Non serve costruire nuove strade (nuovi chip costosi), basta riorganizzare il traffico (il software) per far scorrere tutto più velocemente.

In sintesi:
MXNorm è un modo intelligente per dire: "Non calcolare tutto da capo se hai già le informazioni necessarie nel taschino. Usale per risparmiare tempo e energia, senza rovinare il risultato finale."

È un passo fondamentale per rendere le Intelligenze Artificiali più efficienti, veloci ed economiche da addestrare, permettendo loro di imparare cose ancora più grandi e complesse in meno tempo.

Sommario tecnico

1. Il Problema

L'addestramento di modelli di deep learning su larga scala è stato storicamente limitato dalle prestazioni delle operazioni di moltiplicazione di matrici (matmul). Negli ultimi anni, l'hardware (GPU) ha fatto passi da gigante nel supportare formati numerici a bassa precisione (come FP8 e FP4), migliorando le prestazioni delle moltiplicazioni di matrice fino a 80 volte rispetto alle generazioni precedenti (es. V100 vs GB200).

Tuttavia, c'è un collo di bottiglia asimmetrico:

• Le operazioni di riduzione (come il calcolo della media o della radice quadrata della media quadratica - RMS) e le operazioni elementwise non hanno beneficiato degli stessi miglioramenti di velocità (miglioramenti solo di 5x-9x).
• Questi operazioni sono spesso limitate dalla larghezza di banda della memoria e non possono essere completamente nascoste (overlap) dalle moltiplicazioni di matrice.
• Nelle architetture moderne come i Transformer (es. Llama), il layer di normalizzazione (RMSNorm) è critico per la stabilità dell'addestramento e viene eseguito prima della quantizzazione MXFP (Microscaling Floating Point). Calcolare la RMS su tensori ad alta precisione prima di convertirli in MXFP8 aggiunge un sovraccarico computazionale significativo che non scala bene con l'aumentare della precisione ridotta.

2. Metodologia: MXNorm

Gli autori propongono MXNorm, un'alternativa "drop-in" (sostituzione diretta) a RMSNorm che elimina la necessità di calcolare la RMS separatamente, riutilizzando invece le scale già calcolate per la quantizzazione MXFP.

Concetti Chiave:

• MXFP (Microscaling Floating Point): I tensori vengono divisi in blocchi. Per ogni blocco, viene calcolato un fattore di scala (basato sull'absmax, il valore assoluto massimo) per quantizzare i valori in un formato a bassa precisione (es. E4M3).
• L'Intuizione: Sia la normalizzazione RMS che la quantizzazione MXFP raccolgono statistiche lungo la dimensione nascosta del tensore. Gli autori osservano che se una distribuzione di probabilità viene scalata linearmente, anche il suo absmax e la sua RMS scalano proporzionalmente.
• L'Approssimazione: Invece di calcolare la vera RMS ($\rho$) su tutti gli elementi, MXNorm stima l'inverso della RMS ($\tilde{\rho}$) utilizzando la media generalizzata di ordine $p$ degli absmax dei blocchi ($\tilde{m}$).
  $$\tilde{\rho}_t \approx \tilde{c}(p, B) \cdot \left( \frac{1}{K} \sum_{k=1}^{K} \tilde{m}_{tk}^p \right)^{-1/p}$$
  Dove $K$ è il numero di blocchi, $B$ è la dimensione del blocco e $\tilde{c}$ è un fattore di correzione calcolato tramite simulazioni Monte Carlo.

Implementazione:

1. Fusione: La normalizzazione e la conversione in MXFP (MXCast) vengono fuse in un'unica operazione.
2. Scelta di $p$: Il paper esplora $p=1$ (media aritmetica degli absmax) e $p=2$ (media quadratica).
3. Gestione dei Gradienti: Per l'addestramento, viene utilizzato un straight-through estimator che riutilizza i gradienti di RMSNorm, ma calcolati su una stima ad alta precisione dell'input (ricostruita usando la scala stimata $\tilde{\rho}$ e l'input originale).
4. Fusione con i Pesi: I parametri di guadagno ($\gamma$) della normalizzazione vengono fusi con i pesi del layer lineare successivo, evitando moltiplicazioni elementwise aggiuntive sul tensore quantizzato.

3. Contributi Chiave

• Riduzione dell'Overhead: MXNorm riduce la dimensione della riduzione necessaria per la normalizzazione di un fattore 32x (o più, a seconda della dimensione del blocco), poiché opera su un vettore di scale (uno per blocco) invece che su tutti gli elementi del tensore.
• Teoria della Convergenza: Viene dimostrato teoricamente (Teorema 1) che la media generalizzata degli absmax dei blocchi converge alla RMS reale (a meno di una costante moltiplicativa) all'aumentare del numero di blocchi, assumendo distribuzioni i.i.d.
• Analisi dei Limiti: Viene dimostrato che MXNorm (con $p=2$) mantiene un limite superiore sui valori normalizzati simile a quello di RMSNorm, cruciale per evitare esplosioni dei gradienti, mentre $p=1$ ha limiti superiori più ampi che possono destabilizzare l'addestramento.
• Efficienza Pratica: Implementazione completa che non richiede modifiche architetturali complesse ai modelli LLM esistenti.

4. Risultati Sperimentali

Gli autori hanno validato MXNorm sull'addestramento pre-training di modelli Llama 3 (125M, 1B e 8B parametri) sul dataset SlimPajama.

• Stabilità e Accuratezza:
  • Su modelli piccoli (125M, 1B), MXNorm ($p=1$ e $p=2$) mostra una sensibilità al learning rate simile a RMSNorm.
  • Su modelli grandi (8B), MXNorm con $p=1$ fallisce, mostrando picchi di perdita (loss spikes) dovuti a una stima imprecisa della RMS in presenza di valori anomali (outlier) e a limiti di output troppo ampi.
  • MXNorm con $p=2$ (media quadratica degli absmax) corrisponde perfettamente a RMSNorm, raggiungendo la stessa perdita finale (2.126 vs 2.132) e prestazioni zero-shot comparabili su benchmark OLMES.
• Analisi dei Picchi di Perdita: L'analisi ha rivelato che $p=1$ non riesce a stabilizzare l'addestramento quando emergono feature outlier, mentre $p=2$ fornisce un limite superiore più stretto che previene l'esplosione dei valori.
• Prestazioni (Speedup):
  • Utilizzando torch.compile su hardware NVIDIA GB200, MXNorm offre un speedup fino a 2.4x rispetto a RMSNorm + MXCast isolati.
  • A livello di intero layer Transformer (Llama 3 8B), si ottiene un speedup del 1.3% in MXFP8 e del 2.6% in NVFP4. Sebbene percentualmente piccoli, questi guadagni sono significativi dato che le operazioni di normalizzazione sono un collo di bottiglia crescente man mano che le moltiplicazioni di matrice diventano più veloci.

5. Significato e Impatto

Il lavoro di MXNorm è significativo perché:

1. Riduce il Collo di Bottiglia Non-Matmul: Man mano che l'hardware accelera le moltiplicazioni di matrice, le operazioni di riduzione e normalizzazione diventano il nuovo limite. MXNorm risolve questo problema riutilizzando dati già disponibili.
2. Abilita l'Addestramento Efficiente in Basso Precisione: Permette di utilizzare formati MXFP8/NVFP4 per l'intero flusso di addestramento (inclusa la normalizzazione) senza sacrificare la stabilità o la precisione del modello, semplificando le pipeline di training e inferenza.
3. Scalabilità: Dimostra che è possibile scalare modelli fino a 8B parametri (e potenzialmente oltre) utilizzando questa tecnica, aprendo la strada a modelli ancora più grandi ed efficienti.
4. Generalizzabilità: Il metodo può essere applicato anche a formati di quantizzazione ancora più stretti (es. INT2, ternario) e ad altre forme di quantizzazione basate su blocchi.

In sintesi, MXNorm rappresenta un'ottimizzazione architetturale intelligente che sfrutta la sovrapposizione tra le esigenze di normalizzazione e quantizzazione, trasformando un'operazione costosa in un sottoprodotto gratuito della quantizzazione stessa.