Source-linked AI summary
Quantization-Aware Distillation for NVFP4 Inference Accuracy Recovery
Meng Xin, Sweta Priyadarshi, Jingyu Xin, Bilal Kartal, Aditya Vavre, Asma Kuriparambil Thekkumpate, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Ido Shahaf, Akhiad Bercovich, Kinjal Patel, Suguna Varshini Velury, Chenjie Luo, Zhiyu Cheng, Jenny Chen, Chen-Han Yu, Wei Ping, Oleg Rybakov, Nima Tajbakhsh, Oluwatobi Olabiyi, Dusan Stosic, Di Wu, Song Han, Eric Chung, Sharath Turuvekere Sreenivas, Bryan Catanzaro, Yoshi Suhara, Tijmen Blankevoort, Huizi Mao
TL;DR
NVFP4 quantization can reduce inference cost but may degrade accuracy, and conventional QAT is difficult to apply when LLMs undergo multi-stage post-training or when original data are unavailable. This report uses KL-divergence QAD to distill a full-precision teacher into a quantized student. Across post-trained LLMs and VLMs, QAD consistently recovers near-BF16 accuracy and remains effective with limited-domain data.
Problem
NVFP4 PTQ can cause non-negligible accuracy loss, while QAT is difficult to reproduce for multi-stage SFT, RL, and model-merging pipelines with unavailable or lower-quality training data.
Method
QAD trains an NVFP4-quantized student to match a full-precision teacher using KL divergence rather than task-specific objectives.
Results
QAD reliably recovers NVFP4 models to near-BF16 accuracy across LLM and VLM tasks, including models produced by complex SFT, RL, and model-merging pipelines.
Takeaways & Limitations
QAD is a practical accuracy-recovery approach that remains robust with partial-domain or synthetic data and modest data and compute requirements.
Abstract
from arXiv · showhide
This technical report presents quantization-aware distillation (QAD) and our best practices for recovering accuracy of NVFP4-quantized large language models (LLMs) and vision-language models (VLMs). QAD distills a full-precision teacher model into a quantized student model using a KL divergence loss. While applying distillation to quantized models is not a new idea, we observe key advantages of QAD for today's LLMs: 1. It shows remarkable effectiveness and stability for models trained through multi-stage post-training pipelines, including supervised fine-tuning (SFT), reinforcement learning (RL), and model merging, where traditional quantization-aware training (QAT) suffers from engineering complexity and training instability; 2. It is robust to data quality and coverage, enabling accuracy recovery without full training data. We evaluate QAD across multiple post-trained models including AceReason Nemotron, Nemotron 3 Nano, Nemotron Nano V2, Nemotron Nano V2 VL (VLM), and Llama Nemotron Super v1, showing consistent recovery to near-BF16 accuracy.
1. Introduction
NVFP4 improves efficiency but can lose accuracy on smaller models, while conventional QAT is difficult to apply to modern multi-stage post-training pipelines and incomplete-data settings. The report evaluates QAD as a simpler recovery approach using teacher-student KL divergence and finds strong performance across such models.
- NVFP4 uses 4-bit arithmetic to improve efficiency, but PTQ accuracy losses are often non-negligible for small LLMs.NVFP4 provides higher arithmetic performance and lower memory use than FP8, while its PTQ behavior varies by model size.
- QAT can recover inference accuracy, but modern LLMs make it difficult to replicate multi-stage SFT, RL, and model-merging procedures.The report identifies complex training pipelines as a practical challenge for applying conventional QAT.
- Original training data may be unavailable, and public datasets can be lower quality than the data used to train open models.
- QAD uses a full-precision model as teacher and trains the quantized student with KL divergence rather than task-specific objectives.
- QAD aligns quantized models with high-precision models, remains effective and stable after multi-stage post-training, and tolerates incomplete data coverage.
2. Background and Related Work
NVFP4 combines compact floating-point representation with fine-grained, two-level scaling, while PTQ offers a low-cost alternative to training-based recovery but can fail on small or sensitive models. QAT and distillation provide recovery mechanisms, yet QAD targets the practical difficulties of multi-stage post-training and data access.
- NVFP4 Format: NVFP4 uses 4-bit values, block size 16, per-block E4M3 scales, and a per-tensor FP32 scale for efficient inference and broader dynamic range.Compared with FP8, NVFP4 offers 2-3× higher arithmetic throughput and approximately 1.8× memory reduction.
- Post-Training Quantization: PTQ calibrates quantization parameters from a small representative dataset without training, keeping training and inference decoupled.
- Post-Training Quantization: More advanced PTQ methods reduce quantization error through MSE or KL calibration, clipping, and low-rank compensation techniques.
- Post-Training Quantization: PTQ often struggles on small models and sensitive tasks, while common algorithms may not improve baseline NVFP4 performance.
- Quantization-Aware Training: QAT quantizes weights and activations while retaining high-precision gradients, but it ideally requires the original loss, data, and training procedure.
- Quantization-Aware Training: Quantization-aware reinforcement learning remains less explored for post-hoc accuracy recovery, motivating QAD for RL-trained models.
- Knowledge Distillation: Knowledge distillation transfers teacher information through soft labels, commonly using KL divergence to match probability distributions.
- Related Work: Prior quantized-distillation work has largely focused on CNNs, integer quantization, or specialized LLM methods, whereas this report studies simple KL-based QAD for NVFP4.
3. Quantization-Aware Distillation
QAD trains an NVFP4 quantized student to match a full-precision teacher’s output distribution using KL divergence, rather than reproducing task-specific post-training objectives. Across SFT-heavy and RL-heavy models, it recovers near-BF16 accuracy, outperforms QAT, and remains effective with partial domain data.
- Method Overview: QAD uses KL divergence between the full-precision teacher and quantized student, whereas QAT uses the original task-specific loss.The teacher supplies soft targets representing the original output distribution.
- Method Overview: QAD nearly eliminates divergence from the BF16 teacher’s output distribution, while QAT can match cross-entropy yet substantially alter that distribution.This distinction motivates evaluating distribution alignment in addition to task loss.
- Post-Trained Models: +4.1% on AIME25 and +1.2% on GPQA-D: QAD outperforms QAT for Llama Nemotron Super V1 and recovers near-BF16 performance.On Nemotron Nano 9B V2, QAD also achieves +4.4% on AIME25 and +5.8% on GPQA-D over QAT.
- Post-Trained Models: QAD recovers near-BF16 performance on RL-heavy models, whereas QAT significantly degrades capabilities even when trained with RL-generated samples.Distillation matches the teacher distribution instead of relearning capabilities from cold-start or RL-related data.
- Data Coverage: Math-only or code-only QAD training nearly matches full-data results on AceReason Nemotron, demonstrating cross-domain knowledge transfer.The teacher’s output distributions encode information about domains not explicitly represented in the student’s input data.
- Training and Evaluation Setup: The setup uses model-specific selective quantization and conservative learning rates, while QAD requires substantially less data than the original post-training.Reported convergence amounts range from approximately 0.3B to 6B tokens across the evaluated models.
4. Ablation study
The ablations show that QAD is robust to training-data source and quality, but its learning-rate sensitivity depends on the model’s post-training history. KL divergence is preferred over MSE, and the original BF16 teacher is more effective than a larger teacher.
- Training data quality: QAD remains robust across SFT, RL-generated, mixed, and even random-token data sources when the data amount is held constant.Synthetic BF16-generated data achieves near-BF16 performance, while random tokens remain comparable to the PTQ baseline.
- Learning rate: For SFT-trained models, performance is best at or below the original post-training learning rate, whereas higher rates can degrade performance or destabilize training.Nemotron Nano 9B V2 is optimal at 1e-6; Nemotron Nano 12B v2 VL is optimal at 2e-6, 10× below its original 2e-5 rate.
- Learning rate: RL-trained models benefit from larger QAD learning rates: AceReason Nemotron performs best at 1e-5, above typical RL rates of approximately 1e-6.The recommended overall learning-rate range is 1e-5 to 1e-6, but the optimal value depends on the original training process.
- Additional choices: KL divergence consistently outperforms MSE because it directly measures differences between teacher and student probability distributions.The report uses KL divergence as the standard distillation loss for matching output distributions.
- Additional choices: Using the original BF16 model as teacher outperforms using a larger same-family teacher for efficient accuracy recovery.The report attributes this to the distribution mismatch introduced by a different teacher, which may require more training data.
5. Conclusion
The report presents QAD as a practical method for recovering NVFP4 inference accuracy across LLMs and VLMs, including models produced by complex post-training pipelines. Its robustness to partial or synthetic data supports practical deployment when PTQ alone is insufficient.
- Conclusion: QAD reliably recovers NVFP4 models to near-BF16 accuracy across LLMs and VLMs, including models trained with SFT, RL, and model merging.The evaluated models include Nemotron Nano, Nemotron Nano VL, Llama Nemotron Super, and AceReason Nemotron.
- Conclusion: QAD remains effective with partial-domain, synthetic, or random-token data, reducing the data and compute requirements of accuracy recovery.The report presents these properties as making QAD a practical default when PTQ alone is insufficient.
A. Llama Nemotron VL results
For Nemotron Nano 12B v2 VL, a single-SFT-stage VLM, QAT achieves accuracy comparable to QAD. The accompanying table evaluates both methods across six vision-language benchmarks.
- Llama Nemotron VL results: QAT achieves comparable accuracy to QAD for Nemotron Nano 12B v2 VL, which undergoes only a single SFT stage after pre-training.This differs from the report’s multi-stage post-training cases where QAD is favored.
- Llama Nemotron VL results: Table 10 compares methods on AI2D, ChartQA, DocVQA, InfoVQA, OCRBench, and TextVQA.These are the six listed evaluation benchmarks for the VLM.
B. Nemotron 3 Nano Data Quality Ablation
Nemotron 3 Nano achieves comparable QAD performance across cold-start SFT data, RL-generated data, and their mixture. The SFT+RL mixture performs slightly better.
- Nemotron 3 Nano Data Quality Ablation: All three Nemotron 3 Nano data sources achieve similar performance, with the SFT+RL mixture performing slightly better.The tested sources are cold-start SFT data, BF16-generated data from RL prompts, and a mixture of both.
C. PTQ for Large Models
Large LLMs are empirically more robust to NVFP4 post-training quantization, achieving near-original accuracy without fine-tuning.
- Larger LLMs are more robust to quantization than smaller models.
- Large models achieve near-original accuracy with NVFP4 PTQ without fine-tuning.The cited table describes models with hundreds of billions of parameters.
D. Comparison of QAT/QAD with native quantized training
QAT and QAD quantize weights and/or activations while keeping gradients in high precision, unlike native quantized training, which targets all three GEMMs. Consequently, QAT and QAD do not reduce backward-pass GEMM precision in the same way as native quantized training.
- Native quantized training quantizes forward propagation, weight-gradient, and data-gradient GEMMs.
- QAT and QAD quantize weights and/or activations while leaving gradients in high precision.
- QAT and QAD leave backward-pass weight-gradient and data-gradient GEMMs unquantized.