Technische Samenvatting: TetraJet-v2

1. Probleemstelling

Het vooraf trainen van Large Language Models (LLM) wordt steeds onbetaalbaarder, wat de noodzaak voor Fully Quantized Training (FQT) met lage precisie drijft. Hoewel 4-bit formaten zoals NVFP4 aanzienlijke efficiëntiewinst bieden ten opzichte van FP8 en BF16, blijft het bereiken van stabiele, bijna verliesloze training op deze precisie een kritieke uitdaging. Eerdere pogingen tot volledig 4-bit training hebben geleden onder ernstige prestatiedegradatie.

De auteurs identificeren twee primaire optimalisatieknelpunten die volledig FP4-training belemmeren, die verschillen van eenvoudige kwantisatiefouten:

1. Weight Oscillation (Trilling van gewichten): Een dominante faalmodus waarbij kleine updates aan master-gewichten met hoge precisie ervoor zorgen dat hun gekwantiseerde representaties herhaaldelijk flippen tussen aangrenzende bakken (bijvoorbeeld tussen 0 en 0,5). Dit fenomeen verstoort de optimalisatie, vooral naarmate de leersnelheid afneemt, en verhindert dat het model convergeert naar een stabiel lokaal minimum.
2. Outlier Features (Uitbijterkenmerken): Een klein aantal activatiekanalen vertoont enorme groottes, waardoor dynamische bereiken ontstaan die FP4 niet nauwkeurig kan representeren. Dit leidt tot hoge schalingsfactoren die de precisie van kleinere waarden binnen dezelfde kwantisatiegroep vernietigen.

Bestaande methoden, zoals het NVFP4-recept van NVIDIA, vertrouwen vaak op gemengde precisie (waarbij sommige blokken in BF16 blijven) of slagen er niet in om deze optimalisatiedynamiek volledig aan te pakken in een setting voor vooraf trainen vanaf nul.

2. Methodologie: TetraJet-v2

TetraJet-v2 is een end-to-end 4-bit FQT-methode die specifiek is ontworpen voor NVFP4, waarbij het formaat wordt toegepast op activaties, gewichten en gradiënten in alle lineaire lagen. De methodologie adresseert de geïdentificeerde knelpunten via drie kerncomponenten:

A. Ongewogen Dubbel-Blok Kwantisatie voor NVFP4 Lineaire Lagen

Om de nauwkeurigheid te maximaliseren binnen de numerieke beperkingen van NVFP4 (dat groepen van 16 elementen gebruikt met 8-bit schaling), stellen de auteurs een dubbel-blok kwantisatieschema voor:

• Structuur: Een buitenste blok van 128 elementen wordt globaal geschaald, terwijl binnenste blokken van 16 elementen lokaal worden geschaald. Dit voorkomt het berekenen van de maximale absolute waarde over de hele tensor, wat de hardwarevriendelijkheid en nauwkeurigheid verbetert ten opzichte van per-tensor schaling.
• Ongewogen Gradiëntenschatting: Het ontwerp van de lineaire laag zorgt voor ongewogen gradiëntenschatting door:
  • Tensors in de backward pass uit te lijnen (het schatten van $\hat{X}$ in plaats van $X$).
  • Stochastisch Ronden toe te passen in de backward pass om te waarborgen dat $E[\text{round}(x)] = x$, waardoor bias in gradiëntverwachtingen wordt voorkomen.
  • Een gewichtsgroepsvorm van $1 \times 16$ (niet-vierkant) aan te nemen voor hogere nauwkeurigheid, wat specifieke uitlijning vereist voor de getransponeerde gewichtsmatrix.

B. OsciReset: Onderdrukking van Trillingen

Om trilling van gewichten aan te pakken, introduceren de auteurs OsciReset, het eerste algoritme dat specifiek is ontworpen om dit knelpunt te onderdrukken bij het vooraf trainen van LLM's.

• Detectie: De methode volgt de optimalisatietrajectorie van gewichten en berekent een "Oscillation Risk"-metriek, gedefinieerd als de verhouding tussen de afgelegde afstand door het gekwantiseerde gewicht ($dist_Q$) en de afgelegde afstand door het master-gewicht ($dist_M$). Een hoog risico geeft aan dat een gewicht bakken flippt ondanks kleine updates aan het master-gewicht.
• Mechanisme: Wanneer de leersnelheid laag is en het trillingsrisico een drempelwaarde overschrijdt, wordt het master-gewicht van het trillende element teruggezet naar het midden van zijn huidige kwantisatiebak.
• Voordeel: In tegenstelling tot eerdere methoden die gewichten bevriezen of zware regularisatie gebruiken (wat de globale optimalisatie schaadt), stelt OsciReset het gewicht in staat om verder te optimaliseren op een nieuw traject dat minder vatbaar is voor trillingen, zonder de prestaties direct te degraderen.

C. OutControl: Beheer van Uitbijters

Om uitbijters te beheren, hanteert TetraJet-v2 een hybride strategie:

• Random Hadamard Transformatie (RHT): Toegepast op de berekeningen van de backward pass ($dX$ en $dW$) om de energie van uitbijters gelijkmatiger over kanalen te verdelen, waardoor de kwantisatievariantie wordt verminderd.
• Statische Uitbijterbehoud: Er wordt erkend dat uitbijterkanalen structureel consistent zijn (specifieke kanalen hebben altijd hoge normen); de methode identificeert de top $p\%$ (bijvoorbeeld 10%) van uitbijterkanalen tijdens een korte kalibratiefase. Deze kanalen worden behouden in hogere precisie (bijvoorbeeld FP8) voor zowel forward als backward passes, terwijl de rest van het netwerk in FP4 blijft. Deze statische selectie vermijdt de rekenkundige overhead van dynamische aanpassingen per batch.

3. Belangrijkste Bijdragen

1. Ongewogen Dubbel-Blok Kwantisatie: Een nieuw ontwerp voor NVFP4 lineaire lagen met ongewogen gradiëntenschatting en optimale blokvormen ($1 \times 16$) voor LLM's, die prestaties levert die beter zijn dan bestaande NVFP4-recepten.
2. OsciReset: Het eerste effectieve algoritme om trilling van gewichten in LLM's tijdens vooraf trainen te onderdrukken, waardoor een primaire bron van instabiliteit in training met lage bitbreedte wordt opgelost.
3. OutControl: Een algoritme met gemengde precisie dat RHT combineert voor de backward pass met statisch behoud van uitbijters voor zowel forward als backward passes om nauwkeurigheid te behouden.
4. Uitgebreide Evaluatie: Uitgebreide experimenten voor vooraf trainen op OLMo-2-modellen (tot 370M parameters, 212B tokens) die de effectiviteit van de volledige stack aantonen.

4. Experimentele Resultaten

De auteurs hebben TetraJet-v2 geëvalueerd op OLMo-2-modellen (70M, 150M en 370M parameters) die vanaf nul zijn getraind.

• Vermindering Prestatiekloof: TetraJet-v2-full (met alle voorgestelde methoden) verkleinde de prestatiekloof met training in volledige precisie (BF16) met gemiddeld 51,3% over modelgroottes.
• Perplexiteit (PPL): Op het 370M-model getraind met 212B tokens bereikte TetraJet-v2-full een validatie-PPL van 25,11, vergeleken met 23,70 voor BF16 en 26,20 voor de NVFP4-basislijn van NVIDIA.
• Downstream Taken: De methode behaalde de hoogste gemiddelde prestatie onder alle FP4-methoden op een reeks benchmarks (ARC, BoolQ, MMLU, enz.), waarbij een 370M-model een gemiddelde nauwkeurigheid van 43,60% bereikte.
• Efficiëntie: CUDA-kernel-implementaties bereikten een end-to-end snelheidswinst van 1,67× ten opzichte van FP8 (TransformerEngine) op RTX 5090-hardware, wat praktische prestatiewinst verifieert ondanks de toegevoegde complexiteit van onderdrukking van trillingen en beheer van uitbijters.
• Ablatiestudies:
  • Het verwijderen van OsciReset leidde tot ernstige prestatiedegradatie en een herstel van de loss.
  • OutControl verbeterde de nauwkeurigheid in zowel forward als backward passes.
  • De methode bleek robuust te zijn voor keuzes van hyperparameters (bijvoorbeeld trillingsdrempel, starttijd).

5. Betekenis en Beweringen

Het artikel beweert dat TetraJet-v2 een belangrijke stap voorwaarts is om volledig vooraf trainen van LLM's met lage bitbreedte (4-bit) haalbaar en concurrerend te maken met training in hoge precisie.

• Optimalisatie versus Representatie: De auteurs benadrukken dat de prestatiekloof bij training met lage bitbreedte niet uitsluitend te wijten is aan kwantisatiefouten (representatie), maar kritiek wordt gedreven door optimalisatieproblemen (trilling en uitbijters) die eerdere methoden niet systematisch hebben aangepakt.
• Generaliseerbaarheid: In tegenstelling tot eerdere technieken voor het mitigeren van trillingen die zijn ontworpen voor Vision Transformers (die vaak vertrouwen op fine-tuning of specifieke hyperparameter-tuning), wordt OsciReset gepresenteerd als een eenvoudige, effectieve methode die toepasbaar is op vooraf trainen van LLM's vanaf nul met minimale tuning.
• Praktijktoepasbaarheid: Door een snelheidswinst van 1,67× ten opzichte van FP8 te bereiken en de nauwkeurigheidkloof met BF16 met meer dan 50% te verkleinen, toont het werk aan dat volledig gekwantiseerde training zowel rekenkundige efficiëntie als modelkwaliteit kan bieden, wat potentieel de hardware- en energiekosten van onderzoek naar foundation-modellen kan verminderen.

De auteurs erkennen beperkingen met betrekking tot de schaal van hun experimenten (tot 370M parameters) en het huidige gebrek aan mainstream hardware-ondersteuning voor strategieën met gemengde precisie zoals FP6 × FP4, wat ze verkenden als een potentiële toekomstrichting voor verdere nauwkeurigheidsverbeteringen.