Report KL divergence from the unquantized model in hf_ptq

[cjluo-nv]

What does this PR do?

Type of change: new feature

Adds --kl_divergence to examples/hf_ptq/hf_ptq.py, plus a way to keep and build its references offline. Right after quantization, before export, it reports how far the quantized model has drifted from the unquantized one:

• the mean KL divergence from the unquantized model's next-token distribution;
• both perplexities;
• how often the most likely next token is unchanged.

Scoring copies llama.cpp's llama-perplexity --kl-divergence: consecutive 512-token chunks, 100 by default, with only the second half of each chunk scored. BOS replaces the first token only if the tokenizer prepends one. With the default wikitext2 data (reproduced byte for byte as llama.cpp's wiki.test.raw), the numbers are directly comparable with llama.cpp's for GGUF quants of the same model.

Why: we're adding GGML IQ formats, and next imatrix and block-GPTQ for them. Those need a quality signal that doesn't require exporting a GGUF and running llama.cpp.

Two ways to get the reference.

• In-run (default): the unquantized log-probabilities come from the same model before it is quantized, so it is loaded once. Disabling quantizers afterwards wouldn't work as a substitute, because GPTQ-style algorithms modify weights in place. The reference is held in host memory as float16, chunks × seq_len / 2 × vocabulary values: about 12.7 GB for 100 × 512 tokens of a 248k vocabulary.

• Saved or offline:

  • --kl_divergence_reference ref.pt scores against ref.pt and skips the unquantized pass when the file exists; otherwise it writes ref.pt from this run.
  • python examples/hf_ptq/kl_divergence.py --pyt_ckpt_path <model> --output ref.pt --data <data> builds one without quantizing, using hf_ptq's own model loader.

  References are written with torch.save and read with torch.load(weights_only=True, mmap=True). Scoring against a reference built for another vocabulary raises an error.

Data. --kl_divergence_data takes wikitext2 (default), a UTF-8 text file, or a ModelOpt dataset name such as nemotron-post-training-dataset-v2. Datasets are chat-formatted when both the dataset and the tokenizer support it, so instruct models can be scored on chat-formatted text as well as prose.

Not supported: --use_fsdp2, --low_memory_mode, --specdec_offline_dataset and layerwise.export_dir. The last leaves the model in export form, so it can't be scored.

Usage

# Score in the same run (one model load)
python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path Qwen/Qwen3.8-27B --recipe general/ptq/iq1_s \
  --export_path /tmp/iq1_s --kl_divergence

# Build a reference offline on other data, then reuse it across PTQ runs
python examples/hf_ptq/kl_divergence.py --pyt_ckpt_path Qwen/Qwen3.8-27B --output ref.pt \
  --data nemotron-post-training-dataset-v2
python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path Qwen/Qwen3.8-27B --recipe general/ptq/iq1_s \
  --export_path /tmp/iq1_s --kl_divergence_reference ref.pt

KL divergence vs. the unquantized model over 25500 scored tokens:
  Mean PPL(Q)    :  11.389264 ± 0.182879
  Mean PPL(base) :   6.746027 ± 0.103227
  Mean KLD       :   0.612284 ± 0.005046
  Same top p     :     67.373 ± 0.294 %

Testing

Cross-check against llama.cpp on Qwen3.8-27B. I loaded the BF16 model in transformers and collected the reference with this module: 297,193 tokens, 100 × 255 scored positions, 64 s on one GPU. I then overwrote every quantized weight with the exact dequantized values of two llama.cpp GGUF quants and scored them with kl_divergence. The GGUF-to-HF tensor placement, including the converter's row and column reorders, was recovered by exact BF16 fingerprints against the BF16 GGUF. llama.cpp's numbers are llama-perplexity --kl-divergence on the same files and text:

	This PR (PyTorch)	llama.cpp
BF16 perplexity	6.7460	6.7411
BF16 vs itself: KL / same top	0.000001 / 100 %	—
UD-IQ1_S-layout quant: KL / perplexity / same top	0.6123 ± 0.0050 / 11.389 / 67.37 %	0.6150 ± 0.0050 / 11.390 / 67.07 %
Block-GPTQ IQ quant: KL / perplexity / same top	0.4544 ± 0.0044 / 9.848 / 71.49 %	0.4569 ± 0.0044 / 9.867 / 71.33 %

KL agrees to about 0.4%, within one standard error, and perplexity to 0.2%. The remaining differences are consistent with compute precision: BF16 PyTorch versus llama.cpp's CUDA kernels.

Unit tests. tests/examples/hf_ptq/test_kl_divergence.py (new, CPU, tiny Qwen3) covers:

• chunking and the scored half;
• BOS substitution;
• zero divergence for an unchanged model;
• equality with a direct KL and perplexity computation on a perturbed model;
• the vocabulary-mismatch error and the reference file round trip;
• text-file tokenization without special tokens;
• drawing dataset samples until there are enough tokens.

test_hf_ptq_args.py gains a test that --kl_divergence_reference implies --kl_divergence. 49 tests pass across both files.

Other checks.

• The Hugging Face Salesforce/wikitext wikitext-2-raw-v1 test split, with blank rows written as " \n", equals llama.cpp's wiki.test.raw byte for byte (1,288,556 characters).
• pre-commit is clean on all changed files.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (new opt-in flags)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A (the scoring follows llama.cpp's documented method; no code copied)
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌ (draft)

Additional Information

Overlaps with #2661, which scores BF16-generated continuations with a second model copy. The two could be consolidated, for example by adding #2661's generated-continuation data as another --kl_divergence_data option.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}