Sintesi Tecnica: Ripensare il Bias di Shrinkage nel Pretraining di LLM in FP4

Problematica

Con la scalabilità dei Large Language Models (LLM), i costi computazionali e di memoria del pretraining rendono necessario un passaggio verso formati a precisione inferiore. Sebbene l'FP8 sia stato ampiamente adottato, l'hardware emergente (es. NVIDIA Blackwell/Rubin, AMD MI350) supporta il calcolo nativo in FP4. Le attuali ricette di training FP4, come NVFP4, si basano prevalentemente sul formato E2M1 (2 bit di esponente, 1 bit di mantissa) combinato con la Random Hadamard Transform (RHT) per mitigare la sensibilità agli outlier.

Tuttavia, queste ricette basate su E2M1 soffrono di instabilità del training e di una significativa degradazione della perdita rispetto ai baseline in BF16. Il documento identifica un limite fondamentale, precedentemente trascurato: il Bias di Shrinkage (distorsione da contrazione). Si tratta di un errore di arro rounding negativo sistematico, inerente ai formati di quantizzazione non uniformi come E2M1. L'asimmetria geometrica dei bin rappresentabili di E2M1 causa una costante "contrazione" delle magnitudo sotto l'arrotondamento Round-to-Nearest-Even (RTNE). Crucialmente, il documento sostiene che, sebbene la RHT sia progettata per disperdere gli outlier e migliorare l'utilizzo del codebook, essa esacerba involontariamente questo bias in E2M1, spostando la massa dei tensori nelle regioni di arrotondamento più asimmetriche del formato, portando a un decadimento moltiplicativo del segnale attraverso le reti profonde.

Metodologia

1. Analisi Teorica del Bias di Shrinkage

Gli autori formalizzano il Bias di Shrinkage come un errore atteso negativo nello spazio della magnitudo normalizzata causato dall'asimmetria geometrica dei bin RTNE.

• Origine Geometrica: In griglie non uniformi come E2M1, la spaziatura tra i valori rappresentabili cambia (es. da 0.5 a 1.0 a 1.5). Per un livello $q_i$, le larghezze dei bin sinistro e destro ($\ell_i$ e $r_i$) sono spesso disuguali. Se $r_i > \ell_i$, l'errore atteso condizionale sotto densità locale uniforme è negativo ($\frac{\ell_i - r_i}{2} < 0$).
• Impatto Sistemico: A differenza del rumore a media zero che si annulla, questo bias agisce come un fattore di attenuazione sistematica ($\alpha < 1$) nelle moltiplicazioni di matrici (GEMM). Attraverso $K$ layer, questi fattori si accumulano moltiplicativamente ($\prod \alpha_k \approx \exp(-\sum \delta_k)$), causando un decadimento sistematico del segnale.
• Interazione con RHT: La RHT trasforma i tensori da "limitati dal range dinamico" a "limitati dalla risoluzione locale". Sebbene ciò migliori l'utilizzo dei bucket, forza i dati nelle regioni di magnitudo intermedia dove l'asimmetria di E2M1 è più pronunciata, amplificando così il bias di shrinkage.

2. La Ricetta UFP4

Sulla base del riscontro che le griglie uniformi aggirano questo errore geometrico, gli autori propongono UFP4 (Uniform FP4), una ricetta di training basata su una griglia uniforme di tipo E1M2/INT4.

• Griglia Uniforme: A differenza di E2M1, la griglia E1M2 mantiene una spaziatura uguale tra i livelli ($\ell_i = r_i$), eliminando la fonte geometrica del Bias di Shrinkage.
• Copertura RHT Completa: Poiché la griglia uniforme è priva di bias, UFP4 applica in sicurezza la RHT a tutti e tre i percorsi di GEMM del training: Forward (FPROP), Data Gradient (DGRAD) e Weight Gradient (WGRAD). Le precedenti ricette E2M1 limitavano la RHT al percorso del gradiente dei pesi per evitare l'accumulo di errori nei percorsi non-leaf.
• Stochastic Rounding (SR): La SR è applicata solo al gradiente upstream ($dY$) per preservare le aspettative del gradiente, mentre viene utilizzato l'RTNE altrove.

Contributi Chiave

1. Identificazione del Bias di Shrinkage: Il documento definisce e dimostra formalmente che i formati non uniformi (E2M1) soffrono intrinsecamente di un errore di arrotondamento negativo sistematico dovuto all'asimmetria geometrica, distinto dagli errori di quantizzazione indotti dalla distribuzione.
2. Spiegazione Unificata dell'Instabilità: Fornisce una spiegazione teorica ed empirica per l'instabilità del training nelle esistenti ricette basate su E2M1, mostrando come il Bias di Shrinkage si accumuli moltiplicativamente ed sia esacerbato dalla RHT.
3. Proposta di UFP4: Una nuova ricetta di training a 4 bit che utilizza griglie uniformi (E1M2/INT4), consentendo una RHT pervasiva su tutte le GEMM del training, convertendo il miglioramento dell'utilizzo dei bucket in una maggiore qualità di quantizzazione.
4. Validazione su Scala Industriale: L'approccio è validato su modelli Dense 1.5B, MoE 7.9B e MoE 124B, dimostrando che le griglie uniformi sono formati pratici per il training FP4 di prima classe.

Risultati Sperimentali

Diagnostica di Quantizzazione e GEMM

• Effetto RHT: Su tensori ricchi di outlier, la RHT migliora significativamente il rapporto di bucket effettivo per entrambi i formati. Tuttavia, il relativo Signal-to-Quantization-Noise Ratio (SQNR) diverge: l'E1M2 mostra un SQNR migliorato post-RHT, mentre l'E2M1 spesso degrada ( $\Delta$SQNR negativo) perché i dati spostati finiscono in bin asimmetrici.
• Output GEMM: Nei test su singola GEMM, l'E1M2 preserva o migliora l'SQNR dell'output attraverso i percorsi forward e backward con RHT, mentre l'E2M1 soffre di degradazione, particolarmente nei percorsi backward.

Pretraining End-to-End

• Degradazione della Loss: Nei task di pretraining a lungo termine, UFP4 (E1M2) ottiene costantemente una degradazione della perdita relativa al BF16 inferiore rispetto ai forti baseline E2M1.
  • Dense 1.5B: Errore relativo ridotto da 1.2570% (E2M1) a 0.9673% (UFP4).
  • MoE 7.9B: Ridotto da 2.3596% a 1.8469%.
  • MoE 124B: Ridotto da 1.7308% a 1.3863%.
• Leggi di Scalabilità (Scaling Laws): Il vantaggio di UFP4 persiste attraverso diverse scale di modello (da 10M a 324M parametri). Le curve di perdita adattate mostrano che il gap di E1M2 rispetto a BF16 rimane minore rispetto al gap di E2M1, e la penalità FP4 non cresce con la scala computazionale.
• Studi di Ablazione:
  • Ambito RHT: La copertura RHT completa (tutte e tre le GEMM) produce i risultati migliori per UFP4, mentre limitare la RHT (come nelle ricette E2M1) ne limita le prestazioni.
  • Restrizione del Range: I tentativi di emulare una griglia uniforme limitando il range dinamico di E2M1 (es. cap del valore massimo) non sono riusciti a eguagliare le prestazioni di una E1M2 nativa, sacrificando il range dinamico senza eliminare completamente il bias.

Efficienza

• Fusione dei Kernel: Gli autori dimostrano che la RHT può essere fusa efficientemente con la quantizzazione FP4. Su hardware testato (SM90/SM100), la latenza della fusione RHT+quantizzazione è solo ~1.06x–1.07x rispetto alla quantizzazione standalone, contro ~1.4x–1.6x per le implementazioni non fuse.

Significato e Rivendicazioni

Il documento sostiene che l'attuale standard industriale che tratta E2M1 come l'unico formato FP4 vitale sia subottimale a causa dei suoi limiti geometrici. La significatività di questo lavoro risiede nel:

1. Sfida allo Status Quo: Suggerisce che i futuri acceleratori ML dovrebbero supportare griglie a 4 bit uniformi di tipo E1M2/INT4 come primitive di training di prima classe accanto a E2M1, piuttosto che affidarsi esclusivamente a E2M1.
2. Mismatch tra Formato e Regime: Evidenzia come l'efficacia delle tecniche di mitigazione degli outlier come la RHT dipenda dalla geometria della griglia sottostante. La RHT è benefica per le griglie uniformi ma deleteria per quelle non uniformi in regimi di training profondi.
3. Ricetta Pratica: UFP4 offre una ricetta concreta e scalabile che riduce il divario tra il training FP4 e BF16 senza richiedere modifiche complesse agli stimatori, semplicemente allineando la geometria della griglia con la distribuzione dei tensori post-RHT.

Gli autori concludono che, sebbene E2M1 rimanga adatto per carichi di lavoro limitati dal range e per l'inferenza, l'hardware per il training dovrebbe evolversi per supportare le griglie uniformi al fine di realizzare appieno il potenziale del pretraining FP4.