Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

This work addresses the challenge of recovering inference accuracy for NVFP4-quantized LLMs and VLMs after post-training by introducing quantization-aware distillation (QAD). QAD trains a quantized NVFP4 student to match the output distribution of a full-precision BF16 teacher using a KL divergence loss, offering stability in multi-stage pipelines (SFT, RL, model merging) and robustness to incomplete data coverage. Across multiple models and tasks, QAD consistently recovers near-BF16 accuracy where PTQ or QAT falter, requiring less data and reducing engineering complexity. The results suggest QAD as a practical default for NVFP4 accuracy recovery, with broad applicability to RL-heavy and multi-domain models, and resources to reproduce the findings are provided.

Abstract

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.

Quantization-Aware Distillation for NVFP4 Inference Accuracy Recovery

TL;DR

This work addresses the challenge of recovering inference accuracy for NVFP4-quantized LLMs and VLMs after post-training by introducing quantization-aware distillation (QAD). QAD trains a quantized NVFP4 student to match the output distribution of a full-precision BF16 teacher using a KL divergence loss, offering stability in multi-stage pipelines (SFT, RL, model merging) and robustness to incomplete data coverage. Across multiple models and tasks, QAD consistently recovers near-BF16 accuracy where PTQ or QAT falter, requiring less data and reducing engineering complexity. The results suggest QAD as a practical default for NVFP4 accuracy recovery, with broad applicability to RL-heavy and multi-domain models, and resources to reproduce the findings are provided.

Abstract

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.
Paper Structure (19 sections, 1 equation, 2 figures, 12 tables)

This paper contains 19 sections, 1 equation, 2 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. NVFP4 Format and Post-Training Quantization
  4. Quantization-aware training
  5. Knowledge Distillation for Quantization
  6. Quantization-Aware Distillation
  7. Method Overview
  8. QAD for Post-Trained Models
  9. Robustness to Incomplete Domain Coverage
  10. Training and Evaluation Setup
  11. Ablation study
  12. Training Data Quality
  13. Learning Rate
  14. Additional Choices
  15. Conclusion
  16. ...and 4 more sections

Figures (2)

  • Figure 1: Comparison of quantization-aware training (QAT) and quantization-aware distillation (QAD). QAT trains with next-token prediction (cross-entropy) on target datasets, while QAD uses distillation loss (KL divergence) with the full-precision teacher model providing soft targets.
  • Figure 2: Comparison of quantization-aware training (QAT) and native quantized training. QAT only quantizes the forward pass for inference recovery, while native quantized training quantizes all three GEMMs (Fprop, Wgrad, Dgrad) to reduce training cost. QAD has a similar compute graph as QAT.