FP8 Training con NVIDIA Transformer Engine: dimezza la memoria mantenendo la qualità

I numeri che rappresentano i pesi di un modello AI possono essere memorizzati con diverse quantità di precisione, proprio come la differenza tra scrivere "3.14159265" e "3.1". Più precisione usa più memoria, ma troppo poca precisione degrada la qualità del modello.

Per anni il training ha usato BF16 — 16 bit per numero — come il compromesso standard. NVIDIA ha ora portato questo a FP8 — solo 8 bit per numero — che usa esattamente la metà della memoria. Il problema è che FP8 ha un range numerico molto limitato, e i valori durante il training possono facilmente traboccare o diventare zero se non gestiti attentamente.

Transformer Engine risolve questo problema automaticamente. Per ogni operazione critica (moltiplicazioni di matrici nell'attenzione e nel feed-forward), misura dinamicamente la scala dei valori e applica un fattore di scaling diverso per ciascun tensore, mantenendo i numeri sempre nel range utilizzabile di FP8. Il risultato è che ottieni tutti i vantaggi della memoria ridotta — modelli più grandi sullo stesso hardware, batch più grandi, training più veloce — senza dover scrivere nemmeno una riga di codice in più, e con una differenza di qualità quasi impercettibile.