Technische Samenvatting: QUADS - Stabiliseren van NVFP4 Reinforcement Learning voor MoE via Kwantiseringsfout-alignement

1. Probleemstelling

Reinforcement Learning (RL) voor Mixture-of-Experts (MoE) Large Language Models (LLMs) staat voor een kritieke bottleneck: de generatie van rollouts consumeert vaak meer dan 70% van de totale wall-clock tijd. Om dit te versnellen, hebben onderzoekers naar low-precision inferentie gekeken, waarbij FP8 dient als een stabiele baseline die doorvoer en de mismatch tussen training en inferentie in evenwicht brengt. De Blackwell-architectuur van NVIDIA introduceert echter NVFP4, een natief W4A4-formaat (4-bit gewichten, 4-bit activaties) dat tot wel ~4× de GEMM-doorvoer van BF16 en ~2× die van FP8 biedt.

Het onderzoek identificeert dat het direct toepassen van NVFP4 W4A4 rollout op MoE RL met BF16-training leidt tot onmiddellijke instabiliteit. In experimenten stort naive NVFP4 RL na ongeveer 150 trainingsstappen in, wat wordt gekenmerkt door een snel groeiende kloof tussen de log-waarschijnlijkheden van de rollout en de trainer, wat leidt tot gecorrumpeerde belangrijkheidsratio's (importance ratios) en reward-instorting. De kernuitdaging is dat het grove E2M1-raster van FP4 (slechts 8 positieve representeerbare niveaus per teken) de numerieke bias ernstiger versterkt dan FP8, waardoor standaard correcties voor importance sampling onvoldoende zijn.

2. Methodologie: QUADS

De auteurs stellen QUantization-error Alignment across Dual Sides (QUADS) voor, een framework ontworpen om NVFP4 RL te stabiliseren door de specifieke bronnen van training-inferentie mismatch aan te pakken. De methodologie is gebaseerd op een gedetailleerde foutanalyse die de log-waarschijnlijkheidskloof ($\delta_t$) deelt in gewichtsfout ($\delta^{weight}_t$), activatiefout ($\delta^{act}_t$) en numerieke padverschillen ($\delta^{num}_t$).

2.1 Foutanalyse en Diagnose

Door middel van gecontroleerde ablaties stellen de auteurs vast dat activatiekwantisering de dominante bron van instabiliteit is, en niet gewichtskwantisering.

• Gewichten: Kunnen worden gesynchroniseerd tussen de inferentie- en de training-engines. Omdat gewichten worden gerepliceerd, elimineert het toepassen van hetzelfde kwantiserings-dekwantiseringspad (QDQ) aan de trainer-zijde effectief de gewicht-geïnduceerde mismatch.
• Activaties: Worden online opnieuw berekend tijdens inferentie. Het grove E2M1-raster introduceert een aanzienlijke kwantiseringsfout ($\Delta X_q$) die verder wordt verergerd door inherente engine-drift ($\eta$) tussen de inferentie- en de training-kernels. Symmetrische W4A4 Quantization-Aware Training (QAT), die zowel gewichten als activaties op de trainer fake-kwantiseert, faalt omdat het de drift eerder versterkt dan vermindert.

2.2 Dual-Side Alignment Strategie

QUADS pakt deze asymmetrie aan via twee complementaire componenten:

A. Asymmetrische Quantization-Aware Training (Trainer-zijde)
In plaats van symmetrische W4A4 QAT, stellen de auteurs Asymmetrische W4A16 QAT voor:

• Gewichten: Fake-gekwantiseerd naar NVFP4 (W4) aan de trainer-zijde om het kwantiseringspad af te stemmen op de rollout-engine.
• Activaties: Gehouden in BF16 (A16) aan de trainer-zijde.
• Rationale: Dit voorkomt dat de trainer extra activatiekwantiseringsfouten introduceert die niet kunnen worden uitgelijnd met de online inferentie-activaties. De optimizer werkt de BF16 master-gewichten bij, maar elke forward pass ziet de gediscretiseerde gewichten die door de rollout worden gebruikt, waardoor $\delta^{weight}_t$ wordt geëlimineerd zonder $\delta^{act}_t$ te versterken.

B. Residual Activation Compensatie (Rollout-zijde)
Om de resterende activatiefout ($\delta^{act}_t$) aan te pakken die niet door de trainer kan worden opgelost, gebruikt de rollout-engine een gerichte correctie:

• Mechanisme: Het systeem identificeert "high-residual" kanalen waar de FP4-kwantiseringsfout het grootst is. Het past een tweede-pass kwantisering toe op deze specifieke kanalen, waarbij een residual-correctieterm ($\Delta X_q$) wordt berekend en toegevoegd aan de gekwantiseerde output.
• Implementatie: Een binaire kanaalmasker selecteert de top $k\%$ van de kanalen (bijv. 50%) op basis van de magnitude van de residual. Deze correctie wordt geïmplementeerd via gefuseerde Triton-kernels om overhead te minimaliseren, waarbij de native W4A4 GEMM-doorvoer voor de meerderheid van de berekeningen behouden blijft.
• Doel: Dit vermindert de afstand tussen de rollout-output en de BF16-referentie, waardoor de log-waarschijnlijkheidskloof wordt gedicht zonder terug te vallen op volledige BF16-inferentie.

3. Belangrijkste Bijdragen

1. Diagnose van NVFP4 RL-falen: Het artikel demonstreert empirisch dat activatiekwantisering, en niet gewichtkwantisering, de primaire drijfveer is van log-waarschijnlijkheidsdivergentie en reward-instorting in naive NVFP4 W4A4 RL.
2. Asymmetrisch QAT-schema: Introductie van een trainings-strategie die gewichten naar FP4 fake-kwantiseert terwijl activaties in BF16 blijven, wat effectief het gewichtspad afstemt zonder de activatie-drift te verergeren.
3. Residual Activation Compensatie: Een techniek aan de rollout-zijde die zich richt op kanalen met hoge fouten met een tweede-pass correctie, waardoor de native W4A4-doorvoer behouden blijft terwijl de activatie-geïnduceerde mismatch wordt gemitigeerd.
4. Empirische Validatie: Demonstratie dat de gecombineerde QUADS-pipeline stabiele, full-parameter MoE GRPO-training mogelijk maakt op NVFP4-hardware, met een nauwkeurigheid die gelijk is aan BF16.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op een 25B-parameter MoE-model (2.8B actieve experts) met gebruik van Blackwell GPU's, getraind met GRPO op taken voor wiskundige redenering en codegeneratie.

• Nauwkeurigheid: De QUADS-pipeline behaalde een gemiddelde pass@1 van 72,86% over vier held-out benchmarks (AIME 2024/25, HMMT 2025, LiveCodeBench). Dit komt overeen met de 73,15% behaald door de volledige BF16-baseline en presteert significant beter dan naive NVFP4 RL (51,37%) en FP8 RL (72,88%).
• Stabiliteit: Terwijl naive NVFP4 RL instortte na ~150 stappen, behield QUADS stabiele trainingscurves die vergelijkbaar zijn met BF16. De maximale per-stap log-waarschijnlijkheidskloof werd verminderd van ~1,3 (naive) naar ~0,86 (QUADS), wat dicht bij de onvermijdelijke engine-drift vloer van ~0,74 ligt die in BF16 wordt waargenomen.
• Doorvoer: QUADS leverde ongeveer 16% hogere rollout-doorvoer vergeleken met FP8. Hoewel de residual compensatie een overhead van ~5–9% toevoegde ten opzichte van plain W4A4, bleef de netto versnelling ten opzien van FP8 aanzienlijk.
• Ablatie studies:
  • Het verwijderen van de residual compensatie (alleen W4A16 QAT) stabiliseerde de training, maar liet een merkbaar nauwkeurigheidsverschil zien vergeleken met BF16.
  • Symmetrische W4A4 QAT (het kwantiseren van zowel gewichten als activaties op de trainer) resulteerde in grotere mismatches dan de asymmetrische aanpak, wat de noodzaak bevestigt om trainer-activaties in BF16 te houden.

5. Betekenis en Claims

Het artikel claimt dat QUADS de eerste methode is die stabiele, full-parameter MoE RL-training mogelijk maakt met behulp van native W4A4 FP4 GEMMs op Blackwell-hardware.

De betekenis ligt in het oplossen van het "kwantiseringsfout-alignement" probleem dat specif kind is aan de extreme grofheid van het FP4 E2M1-raster. Door te erkennen dat gewichten en activaties fundamenteel verschillende aligneringsstrategieën vereisen (synchronisatie versus online compensatie), laten de auteurs zien dat het mogelijk is om de enorme doorvoerenwinsten van FP4 te benutten zonder de numerieke stabiliteit op te offeren die vereist is voor policy-gradient optimalisatie. Het werk vestigt dat met de juiste dual-side alignment, NVFP4 een levensvatbaar, hoog-doorvoer alternatief kan zijn voor FP8 voor RL-workloads, waarbij een ~16% doorvoervoordeel wordt geboden terwijl de BF16-niveau modelprestaties behouden blijven.