Technical Summary: QUADS - Stabilizing NVFP4 Reinforcement Learning for MoE via Quantization-error Alignment

1. Problem Statement

Reinforcement Learning (RL) for Mixture-of-Experts (MoE) Large Language Models (LLMs) faces a critical bottleneck: rollout generation often consumes over 70% of the end-to-end wall-clock time. To accelerate this, researchers have explored low-precision inference, with FP8 serving as a stable baseline that balances throughput and training-inference mismatch. However, NVIDIA's Blackwell architecture introduces NVFP4, a native W4A4 (4-bit weights, 4-bit activations) format offering up to ~4× the GEMM throughput of BF16 and ~2× that of FP8.

The paper identifies that directly applying NVFP4 W4A4 rollout to MoE RL with BF16 training causes immediate instability. In experiments, naive NVFP4 RL collapses after approximately 150 training steps, characterized by a rapidly growing gap between rollout and trainer log-probabilities, leading to corrupted importance ratios and reward collapse. The core challenge is that the coarse E2M1 grid of FP4 (only 8 positive representable levels per sign) amplifies numerical bias more severely than FP8, rendering standard importance sampling corrections insufficient.

2. Methodology: QUADS

The authors propose QUantization-error Alignment across Dual Sides (QUADS), a framework designed to stabilize NVFP4 RL by addressing the specific sources of training-inference mismatch. The methodology is grounded in a detailed error analysis decomposing the log-probability gap ($\delta_t$) into weight error ($\delta^{weight}_t$), activation error ($\delta^{act}_t$), and numerical path differences ($\delta^{num}_t$).

2.1 Error Analysis and Diagnosis

Through controlled ablations, the authors determine that activation quantization is the dominant source of instability, not weight quantization.

• Weights: Can be synchronized between the inference and training engines. Since weights are replicated, applying the same quantization-dequantization (QDQ) path on the trainer side effectively eliminates weight-induced mismatch.
• Activations: Are recomputed online during inference. The coarse E2M1 grid introduces significant quantization error ($\Delta X_q$) that is further exacerbated by inherent engine drift ($\eta$) between the inference and training kernels. Symmetric W4A4 Quantization-Aware Training (QAT), which fake-quantizes both weights and activations on the trainer, fails because it amplifies the drift rather than reducing it.

2.2 Dual-Side Alignment Strategy

QUADS addresses this asymmetry through two complementary components:

A. Asymmetric Quantization-Aware Training (Trainer Side)
Instead of symmetric W4A4 QAT, the authors propose Asymmetric W4A16 QAT:

• Weights: Fake-quantized to NVFP4 (W4) on the trainer side to align the quantization path with the rollout engine.
• Activations: Kept in BF16 (A16) on the trainer side.
• Rationale: This prevents the trainer from introducing additional activation quantization error that cannot be aligned with the online inference activations. The optimizer updates BF16 master weights, but every forward pass sees the discretized weights used by the rollout, eliminating $\delta^{weight}_t$ without amplifying $\delta^{act}_t$.

B. Residual Activation Compensation (Rollout Side)
To address the remaining activation error ($\delta^{act}_t$) that cannot be fixed by the trainer, the rollout engine employs a targeted correction:

• Mechanism: The system identifies "high-residual" channels where the FP4 quantization error is largest. It applies a second-pass quantization to these specific channels, computing a residual correction term ($\Delta X_q$) and adding it back to the quantized output.
• Implementation: A binary channel mask selects the top $k\%$ of channels (e.g., 50%) based on residual magnitude. This correction is implemented via fused Triton kernels to minimize overhead, preserving the native W4A4 GEMM throughput for the majority of the computation.
• Goal: This reduces the distance between the rollout output and the BF16 reference, closing the log-probability gap without reverting to full BF16 inference.

3. Key Contributions

1. Diagnosis of NVFP4 RL Failure: The paper empirically demonstrates that activation quantization, rather than weight quantization, is the primary driver of log-probability divergence and reward collapse in naive NVFP4 W4A4 RL.
2. Asymmetric QAT Scheme: Introduction of a training-side strategy that fake-quantizes weights to FP4 while maintaining BF16 activations, effectively aligning the weight path without exacerbating activation drift.
3. Residual Activation Compensation: A rollout-side technique that targets high-error activation channels with a second-pass correction, preserving native W4A4 throughput while mitigating activation-induced mismatch.
4. Empirical Validation: Demonstration that the combined QUADS pipeline enables stable, full-parameter MoE GRPO training on NVFP4 hardware, achieving BF16-level accuracy.

4. Experimental Results

Experiments were conducted on a 25B-parameter MoE model (2.8B active experts) using Blackwell GPUs, trained with GRPO on mathematical reasoning and code generation tasks.

• Accuracy: The QUADS pipeline achieved an average pass@1 of 72.86% across four held-out benchmarks (AIME 2024/25, HMMT 2025, LiveCodeBench). This matches the 73.15% achieved by the full BF16 baseline and significantly outperforms naive NVFP4 RL (51.37%) and FP8 RL (72.88%).
• Stability: While naive NVFP4 RL collapsed after ~150 steps, QUADS maintained stable training curves comparable to BF16. The maximum per-step log-probability gap was reduced from ~1.3 (naive) to ~0.86 (QUADS), approaching the irreducible engine-drift floor of ~0.74 observed in BF16.
• Throughput: QUADS delivered approximately 16% higher rollout throughput compared to FP8. While the residual compensation added a ~5–9% overhead compared to plain W4A4, the net speedup over FP8 remained substantial.
• Ablation Studies:
  • Removing the residual compensation (W4A16 QAT only) stabilized training but left a noticeable accuracy gap compared to BF16.
  • Symmetric W4A4 QAT (quantizing both weights and activations on the trainer) resulted in larger mismatches than the asymmetric approach, confirming the necessity of keeping trainer activations in BF16.

5. Significance and Claims

The paper claims that QUADS is the first method to enable stable, full-parameter MoE RL training using native W4A4 FP4 GEMMs on Blackwell hardware.

The significance lies in resolving the "quantization error alignment" problem specific to the extreme coarseness of the FP4 E2M1 grid. By recognizing that weights and activations require fundamentally different alignment strategies (synchronization vs. online compensation), the authors demonstrate that it is possible to leverage the massive throughput gains of FP4 without sacrificing the numerical stability required for policy-gradient optimization. The work establishes that with the correct dual-side alignment, NVFP4 can serve as a viable, high-throughput alternative to FP8 for RL workloads, offering a ~16% throughput advantage while maintaining BF16-level model performance.