Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[QAD] Request for exact training configurations used in arXiv:2601.20088

Aperta
  • Qualcosa non va?
#2,002 5 commenti 0 reazioni 2 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@mxinO ci sta già lavorando.

Dal 22/7/2026.

Valutazione

Questa issue non è ancora stata valutata.

Descrizione

Question

Thank you for releasing Model Optimizer and the QAD examples. I am trying to understand and reproduce the accuracy recovery reported in Quantization-Aware Distillation for NVFP4 Inference Accuracy Recovery (arXiv:2601.20088v3).

The paper reports near-BF16 recovery for several models. It provides the loss (forward KL), temperature (T=1), approximate token counts, and a few learning rates, but I could not find the exact per-model training configurations used to produce the reported checkpoints.

Could you please share the exact configs/commands (or point me to them) for the paper experiments?

Repository path and scripts I used

I used the Hugging Face/FSDP2 path in examples/llm_qat, not the Megatron-LM path in examples/llm_qad.

Model Optimizer revision:

7ed19b290692f801a54be7266de0c34871a58949
modelopt 0.46.0.dev164+g7ed19b290

The two entry points were:

  1. Pre-quantization: examples/llm_qat/quantize.py
  2. QAD training: examples/llm_qat/train.py, which constructs QADTrainer

The training command was equivalent to:

accelerate launch \
  --config-file configs/accelerate/fsdp2.yaml \
  --num_processes 4 \
  train.py --config /path/to/qad.yaml

The FSDP2 config used BF16 mixed precision, transformer-based wrapping of Qwen3DecoderLayer, activation checkpointing, sharded state dicts, reshard_after_forward: true, and no CPU parameter offload.

Important quantizer caveat

These experiments are not a direct NVFP4 reproduction. I tested ordinary symmetric integer quantization:

  • Model/teacher: Qwen3-8B; the original BF16 Qwen3-8B was also the teacher.
  • Weights: signed symmetric INT4, group size 64, static per-group max scale, narrow range ([-7, 7]).
  • Activations: signed symmetric INT8, group size 64, dynamic per-token/per-group max scale, narrow range ([-127, 127]).
  • KV cache: not quantized.
  • Quantized modules: the 252 decoder GEMM weights (q/k/v/o_proj and gate/up/down_proj); embedding, norms, and lm_head were not quantized.
  • Fake quantization was used in the forward pass with STE in backward.

The ordinary dynamic INT8-g64 fused CUDA fake-quant path was added locally by me because it was not supported by the repository. It computes a fresh symmetric INT8 scale for each contiguous group of 64 activation values. This is not MXINT8, MXFP8, or NVFP4, and I am not claiming that it is an official Model Optimizer quantizer. I mention this explicitly because the different number format may be relevant to the result.

Shared QAD training configuration

Both runs used:

distill: true
criterion: logits_loss
temperature: 1.0
liger_jsd_beta: 0.0       # forward KL
num_train_epochs: 1.0
learning_rate: 1.0e-5
weight_decay: 0.0
lr_scheduler_type: cosine
warmup_ratio: 0.05
per_device_train_batch_size: 2
per_device_eval_batch_size: 2
gradient_accumulation_steps: 2
model_max_length: 8192
use_liger_kernel: true
manual_gc: true
bf16: true
seed: 42

Additional optimizer settings:

  • 4 × A100-SXM4 80 GB, global training batch size = 2 × 4 × 2 = 16
  • adamw_torch_fused, betas (0.9, 0.999), epsilon 1e-8
  • Gradient clipping: 1.0
  • All student parameters remained trainable (trainable_params: null, frozen_params: null), including embedding, norms, and lm_head; the teacher was frozen.
  • Static INT4 amax values were buffers and were not learnable.
  • Loss was forward KL after the causal shift, reduced as the mean over positions where shifted labels != -100.
  • With the chat-template preprocessing in examples/llm_qat/dataset_utils.py, this meant that KL was applied to assistant-response tokens rather than every non-padding token.

Environment:

PyTorch 2.11.0+cu128
Transformers 5.12.1
Accelerate 1.14.0
Datasets 5.0.0
CUDA 12.8

Dataset blend

Both runs used the repository's default QAD/SFT mixture with a 90/5/5 train/eval/test split and the following relative source weights:

Source Subset Ratio
nvidia/Nemotron-SWE-v1 r2e_gym 6000
nvidia/Nemotron-Math-v2 medium 2500
nvidia/Nemotron-Science-v1 MCQ 1500
nvidia/Nemotron-Science-v1 RQA 1500
nvidia/Nemotron-Instruction-Following-Chat-v1 chat_if 5000
nvidia/Nemotron-Post-Training-Dataset-v2 chat 1500
nvidia/Nemotron-Competitive-Programming-v1 competitive_coding_python_part00 1000

Two completed QAD runs

Setting Run 1 Run 2
blend_size 20,000 100,000
train_samples 20,000 100,000
Actual cached training examples 16,580 89,959
Approx. non-padding training tokens 89.3M 457.7M
Approx. assistant/loss tokens 45.2M 233.8M
Optimizer steps 1,036 5,622
Epochs 1 1
eval_samples 2,000 (921 available) 200
eval_steps 50 100
Final validation KL 0.02480 0.02552

Run 2 was a fresh run from the same pre-quantized student, not a continuation of Run 1.

Observed behavior

Training and validation KL decreased normally, and increasing the blend improved held-out output fidelity slightly. However, the final QAD checkpoints still did not outperform my GPTQ/AutoRound PTQ baselines.

For a controlled full-model comparison, I baked the final INT4 fake-quantized values into BF16 tensors, loaded the result as a normal Hugging Face model, and applied the same dynamic symmetric AINT8-g64 fake quantizer to all 252 decoder linear inputs during evaluation.

Forward KL(BF16 teacher || quantized model), each on 512 held-out samples of length 2,048:

Model Pile held-out KL OpenOrca KL
QAD Run 1 0.04901 0.14227
QAD Run 2 0.04678 0.13437
AutoRound best WINT4-g64 + the same AINT8-g64 evaluation quantizer 0.02718 0.09104

For Run 1, five-run downstream evaluation with the same AINT8-g64 activation fake quantizer gave:

Model AIME-2026 GPQA-Diamond Simple average
QAD Run 1 0.6000 0.5778 0.5889
AutoRound best WINT4-g64 0.6400 0.5616 0.6008

I understand that:

  1. the paper evaluates NVFP4 rather than this custom integer W4A8 format;
  2. QAD is not necessarily expected to beat every advanced PTQ algorithm; and
  3. model family, post-training history, data volume, and checkpoint selection can materially change the result.

I am therefore mainly trying to determine whether my training setup differs from the setup that produced the paper's reported recovery.

Details that would help reproduce the paper

Could you share the following for each reported model/checkpoint, ideally as runnable YAML and launch commands?

  1. Which code path was used: examples/llm_qad (Megatron-LM), examples/llm_qat (HF Transformers), or the NeMo path?
  2. Exact optimizer, LR, scheduler, warmup, weight decay, gradient clipping, global/micro batch size, gradient accumulation, sequence length, number of steps/epochs, and distributed configuration.
  3. Exact training-data mixture and preprocessing: packing, truncation, chat template, whether responses were generated by the teacher, and whether samples were reused across epochs.
  4. Was KL reduced over all non-padding tokens or only assistant/response tokens? Was it a global token mean across data-parallel ranks?
  5. Which student parameters were updated: all weights, only quantized GEMM weights, lm_head, embeddings, norms, and/or quantizer scales?
  6. Exact NVFP4 recipe and calibration details: quantized/excluded modules, weight and activation scale behavior during QAD, calibration algorithm/size, KV-cache format, and whether scales were fixed or learnable.
  7. Validation/checkpointing details: validation set size, validation metric/reduction, evaluation interval, save interval, and how the “top 10 lowest-validation-loss checkpoints” were selected and then ranked on downstream benchmarks.
  8. Do you expect the same QAD recipe to improve ordinary symmetric INT4-g64 / INT8-g64 fake quantization, or are there NVFP4-specific considerations that make a different recipe necessary?

Publishing the exact configs used for the paper (even for one model such as AceReason Nemotron 1.1 7B or Nemotron Nano 9B V2) would be extremely helpful.

Thank you.

Lingua principale
Python
Stelle
5.3k
Fork
745
Merge medio
4g 2h
PR unite (30g)
160

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/Model-Optimizer

Tutte le issue di NVIDIA/Model-Optimizer

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.