Add optional MLflow tracking to hf_ptq.py

[cjluo-nv]

What does this PR do?

Type of change: new feature

Adds modelopt.torch.utils.mlflow.MlflowRunLogger, a reusable helper for recording a script run on an MLflow tracking server, and wires examples/hf_ptq/hf_ptq.py up to it via --mlflow <tracking-uri> so a PTQ run can be reproduced from its MLflow entry alone. Without the flag, behavior is unchanged — every hook is gated on it.

The logger lives in the library rather than the example so other scripts can record runs the same way: it takes a tracking URI, an experiment name and an explicit enabled flag, with params, tags and artifacts passed in. hf_ptq.py supplies only the PTQ-specific pieces (its params, the resolved recipe, the quantization summaries). mlflow is an optional dependency, imported only once tracking is enabled, so it is not a new requirement for the library.

The run is opened before the model loads, so a bad URI or an unreachable server fails in seconds rather than after hours of calibration. The invocation and the recipe are uploaded at that point too, which keeps a crashed run useful: it is still recorded, with status FAILED and its log attached.

Uploaded artifacts:

Artifact	Contents
command.txt	The full invocation, copy-pasteable
version.txt	The ModelOpt version that ran (also a searchable tag)
recipe/resolved_recipe.yaml	The --recipe with $imports expanded
logs/hf_ptq.log	Everything the run printed, including a crash traceback
summary/quant_summary.txt	Per-quantizer summary (unless --no-verbose)
summary/moe.html	Per-expert calibration token counts, when the run produces them

Plus model / format / calibration settings as searchable params, and user / hostname / modelopt_version / git_sha tags.

Three design points worth review:

1. The recipe is uploaded resolved, not verbatim. A recipe may be a directory or use $imports, so the source file is not self-contained. For huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe the source is 2,230 B / 58 lines against 7,563 B / 308 lines resolved — the raw file records under 30% of what actually ran.
2. hf_ptq.py has no logging framework (bare print()), so the log is produced by teeing stdout/stderr. Handlers that libraries bound to sys.stderr at import time are re-pointed at the tee for the run's duration and handed back afterwards; without that, transformers / huggingface_hub warnings reach the console but never the log. Native (C-level) output is still not captured — documented in the README.
3. The recipe upload lives in the caller, not the library. That keeps modelopt.recipe out of modelopt.torch.utils, which would otherwise risk a modelopt.torch.utils → modelopt.recipe → modelopt.torch.quantization → modelopt.torch.utils import cycle.
4. MLflow failures never fail the quantization. Startup validation is fatal by design (it is before any GPU work); the end-of-run upload is best-effort.

Only the main rank uploads, so --use_fsdp2 runs produce a single run.

Usage

python hf_ptq.py \
  --pyt_ckpt_path <huggingface_model_card> \
  --recipe general/ptq/nvfp4_default-kv_fp8_cast \
  --export_path <quantized_ckpt_path> \
  --mlflow https://<your-mlflow-server>/

[mlflow] experiment: $USER/hf_ptq/<checkpoint basename>-<recipe name>
[mlflow] run: https://<your-mlflow-server>/#/experiments/13/runs/c243352e...

--mlflow_experiment and --mlflow_run_name override the defaults ($USER/hf_ptq/<basename>-<recipe name or --qformat>, and the UTC start time). Passing --mlflow with no value uses $MLFLOW_TRACKING_URI. Authentication uses MLflow's own env vars.

Testing

Unit — 26 tests in tests/unit/torch/utils/test_mlflow.py for the library, plus 8 in tests/examples/hf_ptq/test_hf_ptq_args.py for the hf_ptq wiring. CPU-only, no network and no mlflow dependency (driven against a stub module). Covers experiment-name derivation and sanitization, URI accept/reject, tee pass-through, the pre-bound-handler redirect, artifact renaming, skipping absent optional outputs, the disabled path, and version.txt. 56 tests pass together with the existing test_hf_ptq_args.py / test_example_utils.py.

Hardware — real PTQ runs against a live MLflow server:

Run	Result
Qwen3-0.6B, NVFP4 PTQ, 1×B200	FINISHED, all artifacts, sane post-quant generations
Qwen3.6-35B-A3B MoE, AutoQuantize w4a16_nvfp4_fp8_at_6p0bits-active_moe, 2×B200	FINISHED in 63 min, search hit effective bits: 6.00; 106 KB log capturing every per-layer decision, 4.4 MB quant summary
Qwen3.6-35B-A3B, plain NVFP4 PTQ, 2×B200	FINISHED
Qwen3-0.6B re-run after the library move, 1×H200	FINISHED, all five artifacts including version.txt
Crash mid-run (gated HF dataset)	FAILED recorded with log + traceback attached, summaries correctly absent
Malformed URI	Rejected by argparse with a Did you mean https://…? hint
Unreachable host	Fails in 9.9 s total, before any model load
No --mlflow	Exit 0, no MLflow output, unchanged export

Coverage gap, stated plainly: summary/moe.html is verified only against a synthetic file (unit test + a real upload). It could not be produced naturally — expert_token_count buffers live on _QuantSparseSequentialMoe, while Qwen3.5/3.6 experts take the fused _QuantFusedExperts path, so no such file is written for these models regardless of --moe_calib_experts_ratio. The uploader's conditional is correct; the branch simply had no natural input available here.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — new optional flags only; no --mlflow means no behavior change.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ — adds mlflow to examples/hf_ptq/requirements.txt. Apache-2.0 (permissive). No code copied from other sources.
• Did you write any new necessary tests?: ✅ — 29 new unit tests.
• Did you update Changelog?: ✅ — 0.47 New Features.
• Did you get Claude approval on this PR?: ❌ — not yet run.

Additional Information

Draft: the summary/moe.html path would benefit from one run on a sparse-sequential MoE model before this is marked ready.

🤖 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

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 10724ec2-8742-4ff6-84d8-ea365d9fd390

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/hf-ptq-mlflow-tracking

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 66.82%. Comparing base (33d05b0) to head (e1b798f).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2023      +/-   ##
==========================================
- Coverage   66.83%   66.82%   -0.02%     
==========================================
  Files         519      519              
  Lines       58916    58916              
==========================================
- Hits        39376    39370       -6     
- Misses      19540    19546       +6     

Flag	Coverage Δ
unit	54.88% <ø> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2023/
Built to branch gh-pages at 2026-07-28 15:13 UTC. Preview will be ready when the GitHub Pages deployment is complete.