Quantization-Aware Distillation for NVFP4 Inference Accuracy Recovery

Dit technische rapport introduceert kwantisatiebewuste distillatie (QAD) als een robuuste en stabiele methode om de nauwkeurigheid van NVFP4-gekwantiseerde grote taalmodellen en visueel-taalmodellen, die via meervoudige post-trainingpijplijnen zijn ontwikkeld, te herstellen tot bijna BF16-niveau zonder volledige trainingsdata.

De kern

Stel je voor dat je een zeer slimme, maar enorme robot hebt die alles kan: wiskunde oplossen, code schrijven, en vragen beantwoorden. Deze robot is zo groot dat hij een hele fabriek aan stroom en ruimte nodig heeft om te werken. Dat is niet praktisch voor een gewone gebruiker.

De onderzoekers van NVIDIA wilden deze robot "verkleinen" zodat hij op een gewone computer of telefoon kan draaien. Ze deden dit door de getallen die de robot gebruikt om te rekenen, af te ronden naar een veel kleinere, minder nauwkeurige vorm (noem het van "gouden munten" naar "kleine koperen muntjes"). Dit noemen ze NVFP4.

Het probleem? Door dit verkleinen verloor de robot een deel van zijn intelligentie. Hij begon fouten te maken, alsof hij zijn bril had verloren.

Hier komt de oplossing uit dit rapport: Quantization-Aware Distillation (QAD). Laten we dit uitleggen met een paar creatieve vergelijkingen.

1. Het probleem: De "Gouden" vs. de "Koperen" Robot

Stel je voor dat je een meesterkok hebt (de BF16-robot, de originele, perfecte versie). Hij kookt een perfecte soep.
Nu wil je een kopie maken die in een klein tentje kan werken. Je gebruikt daarvoor goedkopere potten en lepels (de NVFP4-robot).
Als je de goedkope robot gewoon de soep laat maken, wordt het een ramp. De smaken zijn niet goed, de groenten zijn te groot of te klein. De robot is "verward" door de nieuwe, beperkte gereedschappen.

2. De oude aanpak: "Probeer het zelf" (QAT)

Vroeger probeerden ze de goedkope robot te leren door hem te laten koken op basis van recepten (data) en te zeggen: "Je hebt de soep te zout gemaakt, doe er minder zout bij."
Dit werkt goed als je de originele recepten en de perfecte keuken hebt. Maar bij moderne AI-modellen is dat lastig:

• We hebben vaak niet meer de originele recepten (de trainingsdata is geheim of te groot).
• De robot is al zo complex getraind (met beloningen voor goed gedrag) dat als je hem opnieuw probeert te leren met oude recepten, hij zijn nieuwe vaardigheden vergeet. Het is alsof je een olympisch atleet laat terugkijken naar zijn basisschool-oefeningen; hij raakt in de war en presteert slechter.

3. De nieuwe oplossing: De "Geestelijke Meester" (QAD)

In plaats van de goedkope robot te laten proberen de recepten zelf te begrijpen, laten we hem kijken naar de meesterkok.

• De Meester (Teacher): De originele, perfecte robot. Hij geeft niet alleen het antwoord ("De soep is klaar"), maar ook de gevoelens en nuances ("De soep is 80% zout, 20% kruidig, en voelt warm aan").
• De Leerling (Student): De kleine, goedkope robot.
• De Les: De leerling probeert niet om het antwoord zelf te raden. Hij probeert exact hetzelfde te voelen als de meester. Hij kijkt naar de output van de meester en zegt: "Ah, de meester denkt dit, dus ik moet ook dit denken."

Dit noemen ze Distillatie. Het is alsof je de "geest" van de meester overbrengt naar de leerling, zonder dat de leerling de hele geschiedenis van de meester hoeft te herbeleven.

Waarom werkt dit zo goed?

1. Het werkt ook voor complexe robots: Veel moderne AI-modellen zijn getraind in meerdere stappen (eerst leren, dan belonen, dan samenvoegen). De oude methode faalde hier omdat je die hele complexe geschiedenis niet kunt nabootsen. De nieuwe methode (QAD) zegt: "Het maakt niet uit hoe de meester het heeft geleerd; kijk gewoon naar wat hij nu denkt."
2. Het is niet kieskeurig over de data: Je hoeft geen perfecte receptenboeken te hebben. Zelfs als je de leerling alleen maar laat kijken naar willekeurige zinnen of alleen maar wiskundige voorbeelden, kan hij de "geest" van de meester overnemen. De meester weet immers al alles; de leerling hoeft alleen maar mee te denken.
3. Het herstelt de intelligentie: Door deze methode te gebruiken, komt de kleine, goedkope robot (NVFP4) bijna precies terug naar het niveau van de grote, dure robot (BF16). De fouten die door het verkleinen waren ontstaan, worden "opgeveegd".

De conclusie in het kort

De onderzoekers hebben een nieuwe manier gevonden om slimme AI's kleiner en sneller te maken zonder dat ze dom worden. In plaats van ze te dwingen om alles opnieuw te leren (wat vaak mislukt), laten ze ze kijken naar de perfecte versie van zichzelf en proberen die gedachten na te bootsen.

Het is alsof je een student niet laat studeren uit een verouderd boek, maar hem laat meedenken met een Nobel-prijswinnaar. De student wordt niet alleen slimmer, maar hij wordt ook slim in een klein pakketje, klaar om op elke computer te draaien.

Technische samenvatting

Titel: Quantization-Aware Distillation (QAD) voor NVFP4 Inference Accuracy Recovery

1. Het Probleem

De snelle expansie van Large Language Models (LLMs) en Vision-Language Models (VLMs) vereist efficiëntere numerieke formaten om rekentijd, geheugengebruik en energieverbruik te verlagen. NVIDIA's NVFP4 (een 4-bit floating-point formaat) biedt een aanzienlijke prestatieverbetering (2-3x hogere rekenthroughput) en geheugenreductie (50%) ten opzichte van FP8.

Hoewel Post-Training Quantization (PTQ) goed werkt voor zeer grote modellen, ondervinden kleinere modellen en modellen die zijn getraind via complexe, multi-stap post-training pijplijnen (zoals Supervised Fine-Tuning (SFT), Reinforcement Learning (RL) en model merging) vaak een niet-verwaarloosbare daling in nauwkeurigheid wanneer ze naar NVFP4 worden gekwantiseerd.

Bestaande methoden voor nauwkeurigheidsrecovery hebben beperkingen:

• Quantization-Aware Training (QAT): Vereist het opnieuw doorlopen van de volledige trainingspijplijn met kwantiseringsbewaking. Dit is vaak onuitvoerbaar omdat de originele trainingsdata niet beschikbaar is, de pijplijnen te complex zijn (vooral bij RL), en het risico bestaat dat de getrainde capaciteiten van het model worden beschadigd.
• Data-afhankelijkheid: Traditionele methoden vereisen vaak volledige datasets van hoge kwaliteit, wat een barrière vormt voor open modellen.

2. Methodologie: Quantization-Aware Distillation (QAD)

Het rapport introduceert QAD als een praktische oplossing om de nauwkeurigheid van NVFP4-gekwantiseerde modellen te herstellen zonder de originele trainingspijplijn te hoeven repliceren.

• Kernprincipe: QAD gebruikt het originele full-precision (BF16) model als een leraar (teacher) en traint het gekwantiseerde model als een student.
• Verliesfunctie: In tegenstelling tot QAT, dat gebruikmaakt van taakspecifieke loss-functies (zoals cross-entropy voor next-token prediction), gebruikt QAD KL-divergentie (KL Divergence) om de output-verdeling van de student te laten overeenkomen met die van de leraar.
  • Formule: $\mathcal{L}_{QAD} = D_{KL}(p_{teacher} || p_{student})$
• Data-eisen: QAD is robuust tegen onvolledige data. Het kan werken met:
  • Een subset van de oorspronkelijke SFT-data.
  • Synthetische data gegenereerd door het model zelf (zelfdistillatie).
  • Zelfs willekeurige tokens (hoewel dit minder optimaal is, breekt het het model niet).
• Toepassing op RL-modellen: Voor modellen die zijn getraind met Reinforcement Learning is QAT vaak schadelijk omdat het de getrainde capaciteiten kan vernietigen. QAD omzeilt dit door simpelweg de output-verdeling van de leraar te kopiëren, zonder de complexe RL-dynamiek opnieuw te hoeven simuleren.

3. Belangrijkste Bijdragen

1. Superieure Alignement: QAD aligneert het gekwantiseerde model veel nauwkeuriger met de oorspronkelijke BF16-verdeling dan QAT, wat resulteert in een lagere KL-divergentie en behoud van de modelcapaciteiten.
2. Stabiliteit bij Multi-stage Pipelines: De methode werkt uitzonderlijk stabiel voor modellen die zijn doorlopen via complexe pijplijnen (SFT + RL + Merging), waar QAT vaak instabiel is of faalt.
3. Robuustheid tegen Data: QAD kan nauwkeurigheid herstellen met beperkte data (bijv. alleen wiskunde-data voor een model dat ook code kan) en faciliteert effectieve cross-domain kennisoverdracht.
4. Praktische Implementatie: Het rapport biedt best practices, hyperparameters (zoals leerlijnen tussen 1e-6 en 1e-5) en code voor verschillende frameworks (Megatron-LM, NeMo, HuggingFace).

4. Resultaten

De auteurs hebben QAD geëvalueerd op diverse modellen, waaronder Nemotron Nano, Llama Nemotron Super, en AceReason Nemotron.

• Nauwkeurigheids Herstel: QAD slaagt er consequent in om NVFP4-modellen terug te brengen naar nagenoeg BF16-nauwkeurigheid op benchmarks zoals MATH500, AIME25, GPQA-Diamond en LiveCodeBench.
• Vergelijking QAT vs. QAD:
  • Bij SFT-georiënteerde modellen presteert QAD beter dan QAT, vooral op redeneertaken (bijv. +4.1% op AIME25 voor Llama Nemotron Super V1).
  • Bij RL-georiënteerde modellen faalt QAT vaak (het degradeert de prestaties onder het PTQ-niveau), terwijl QAD succesvol de BF16-prestaties herstelt.
• Data-efficiëntie: QAD vereist aanzienlijk minder data dan de originele post-training (bijv. ~0.3 miljard tokens voor een 49B model versus de oorspronkelijke datasetgrootte).
• Kwaliteit van Data: Experimenten tonen aan dat QAD robuust is; zelfs data gegenereerd door het model zelf (synthetisch) of gefilterde correcte antwoorden werken goed. Interessant genoeg presteert het gebruik van alle gegenereerde data (inclusief fouten) soms beter dan alleen correcte data, omdat fouten ook waardevolle informatie bevatten voor distillatie.

5. Betekenis en Conclusie

Dit rapport positioneert Quantization-Aware Distillation als de de facto standaard voor het herstellen van nauwkeurigheid in NVFP4-inferentie, vooral voor moderne, complex getrainde LLMs en VLMs.

• Praktische Impact: Het lost het dilemma op tussen de wens naar extreem efficiënte inferentie (4-bit) en de noodzaak om modelkwaliteit te behouden. Het maakt NVFP4-deployments haalbaar zonder de enorme kosten en complexiteit van het opnieuw trainen van modellen.
• Toekomstperspectief: QAD biedt een weg om kwantiseringsfouten te minimaliseren zonder toegang te hebben tot de originele trainingsdata of de volledige trainingsinfrastructuur, wat essentieel is voor de bredere adoptie van lage-bit inferentie in de industrie.

Samenvattend biedt QAD een stabiele, data-efficiënte en robuuste methode om de kloof tussen 4-bit kwantisatie en full-precision prestaties te overbruggen, zelfs voor de meest geavanceerde modellen met RL-training.