Quantization-aware distillation yields faster Nemotron 3.5 Lightning NVFP4 checkpoint

NVIDIA applied quantization‑aware distillation (QAD) to the open Nemotron 3.5 Lightning model to produce a deployable NVFP4 checkpoint that reduces memory footprint while recovering accuracy lost to aggressive quantization. The workflow combines post‑training quantization (PTQ) to create an aggressively quantized student checkpoint and a second stage that distills the frozen BF16 teacher into the quantized student using KL divergence on logits. The full recipe is available in NVIDIA Model Optimizer and the final NVFP4 checkpoint has been published on Hugging Face.

What was done and why it matters

• The Nemotron 3.5 Lightning NVFP4 checkpoint can be compressed to roughly 22 GB from a ~66 GB BF16 baseline (detailed evaluations use ~21.19 GB PTQ/QAD checkpoints versus ~65.85 GB BF16), enabling up to ~3–4× throughput improvements at deployment.
• Post‑training quantization (PTQ) alone is a fast way to compress models, but very aggressive quantization (for example taking many weights to 4 bits) introduces accuracy degradation. Quantization‑aware distillation (QAD) lets the quantized student adapt to quantization noise by learning from a frozen BF16 teacher, recovering most of the lost accuracy.

The QAD approach

• Two stages:
  1. PTQ: quantize the BF16 teacher to produce a W4A16 / NVFP4 student checkpoint. The authors intentionally allow PTQ to be aggressive because QAD will later recover accuracy.
  2. QAD: during training the student runs forward passes with simulated quantization; a KL divergence loss on logits trains the student to match the frozen BF16 teacher.

PTQ recipes and calibration

• The team evaluated multiple PTQ recipes differing in calibration method (max/dynamic vs MSE/static), Mamba projection formats, and KV cache quantization. All examined recipes used W4A16 NVFP4 weights but varied in whether scales are dynamic or frozen during QAD.
• The recommendation is to consider more aggressive PTQ when QAD follows (e.g., Mamba linear layers to W4A16 rather than FP8) because QAD can recover the induced accuracy loss.

QAD training configuration (concrete settings)

• Distillation runs both BF16 teacher and NVFP4 student on the same batch; student is trained with KL loss on logits while teacher is kept frozen.
• Example hyperparameters from NVIDIA’s runs:
  • learning rate: constant 5e‑6 (no warmup)
  • dropout: disabled
  • gradient clipping: 1.0
  • hardware example: 2 nodes × 8 GPUs (TP=2, EP=4)
  • short experimental run: train samples 6,400 (≈400 iterations)
  • final runs used much larger effective token counts, e.g., sequence lengths up to ~524k for long‑context preservation
• They recommend starting reproducibility tests from NVIDIA’s released Nemotron‑Post‑Training v1/v2 datasets.

Quantization scale strategies

• Dynamic scale QAD: used with max‑calibrated PTQ checkpoints (scales recomputed during training, weights and scales both adapt).
• Frozen scale QAD: used with MSE‑based PTQ checkpoints (scales found at PTQ time are frozen; only weights update during QAD).
• The PTQ recipe determines which scale strategy to use; trying multiple PTQ variants is advised.

Evaluation results

• Storage footprint: aggressive W4A16 PTQ and QAD checkpoints both measured at 21.19 GB in the reported experiments, compared with a BF16 baseline of 65.85 GB.
• Checkpoint A (SFT intermediate): PTQ reached 96.33% median score recovery vs BF16; after QAD median recovery rose to 99.72% (+3.39). Ten of 11 benchmarks improved under QAD. Example: AIME 2025 dropped 3.70 points under PTQ but after 200 QAD iterations was within 0.57 points of BF16.
• Checkpoint B (RL intermediate): PTQ 95.84% median → QAD 98.53% (+2.69). AA v4.1 Index climbed from 20.03 to 23.48 after QAD (BF16: 24.81).
• Final NVFP4 checkpoint (conservative PTQ tuned for accuracy): PTQ median recovery 99.24%, QAD 98.97% — here PTQ already sat close to BF16, so QAD gains are smaller overall but visible on agentic and coding benchmarks (e.g., Terminal‑Bench v2.1, SWE‑Bench Multilingual, PinchBench, HLE).

Reproducibility and tools

• The full QAD recipe ships in NVIDIA Model Optimizer as megatron_lm_qad.yaml for the end‑to‑end pipeline. A Megatron‑Bridge variant (mbridge_qad.yaml) is provided for PyTorch‑native runs within the NeMo Framework.
• Example launcher commands and modelopt.torch.quantization usage are included in the published materials. The final NVFP4 checkpoint is published on Hugging Face.

Conclusions

QAD enables more aggressive quantization choices at the PTQ stage by recovering the majority of the accuracy lost to low‑precision grids during training, allowing substantially smaller checkpoints (≈21–22 GB) with near‑BF16 performance across many benchmarks. For teams that must balance latency, memory and throughput, the QAD pipeline in NVIDIA Model Optimizer is a reproducible option to push precision/size tradeoffs further while maintaining high quality.

Acknowledgements

NVIDIA thanks the Nemotron team for open‑sourcing Nemotron 3.5 Lightning and the post‑training datasets, and credits contributors to the recipe and evaluations named in the source material (Asma Kuriparambil Thekkumpate, Carlo del Mundo, Chenjie Luo, Daniel Lo, Daria Levy, Frank Sun, Hung‑Yueh Chiang, James Shen, Jinhang Choi, Sweta Priyadarshi, Konstantinos Krommydas, Meng Xin, Rohan Joshi, Wei‑Ming Chen, Victor Cui, Trenton Starkey, Yaniv Galron).