Move the MLflow wiring into example_utils and simplify the flag handling

hf_ptq.py is down from ~90 lines of tracking glue to three calls:
add_mlflow_args(parser), resolve_mlflow_args(args, parser) and
`with mlflow_run(args):`. Everything that knows about the flags, the
params, the tags and the summary paths now sits beside the other hf_ptq
helpers in example_utils.

Simplify how the URI is settled. Provenance is now just "did the user type
the flag", read before the fallback, so the whole thing is one assignment
and one branch:

    args.mlflow_required = args.mlflow is not None
    args.mlflow = args.mlflow or os.environ.get("MLFLOW_TRACKING_URI") or None

Dropping nargs="?" is what makes that work. The bare --mlflow form meant
"use $MLFLOW_TRACKING_URI", which is now what happens with no flag at all,
so it was redundant -- and it was the only source of the empty-string case
that forced the three-way logic. A bare --mlflow now gets argparse's own
"expected one argument".

Move the MLflow section of the README to the end, per review: it is a lot
of detail to meet before the feature examples like AutoQuantize. Added to
the contents table so it stays findable.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

Author: cjluo-nv

examples/hf_ptq/README.md

@@ -19,6 +19,7 @@ This section focuses on Post-training quantization, a technique that reduces mod
 | Evaluate Accuracy | Evaluate your model's accuracy! | \[[Link](#evaluate-accuracy)\] | |
 | Exporting Checkpoints | Export to Hugging Face Unified Checkpoint and deploy on TRT-LLM/vLLM/SGLang | \[[Link](#exporting-checkpoints)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/deployment/3_unified_hf.html)\] |
 | Pre-Quantized Checkpoints | Ready to deploy Hugging Face pre-quantized checkpoints | \[[Link](#pre-quantized-checkpoints)\] | |
+| Tracking runs with MLflow | Record a PTQ run on an MLflow server so it can be reproduced from its entry alone | \[[Link](#tracking-runs-with-mlflow)\] | |
 | Resources | Extra links to relevant resources | \[[Link](#resources)\] | |
 
 </div>
@@ -293,61 +294,6 @@ scripts/huggingface_example.sh --model <model> --quant nvfp4 --vlm --calib_with_
 > Note: when `--calib_with_images` is set, `--calib_size` must be a single value, and the calibration dataset is nvidia/nemotron_vlm_dataset_v2.
 This functionality is currently in beta and has been tested on `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16`.
 
-### Tracking runs with MLflow
-
-Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow <tracking-uri>`, to record a PTQ
-run on an MLflow server so it can be reproduced later from its MLflow entry alone:
-
-```bash
-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>/
-```
-
-The run is opened *before* the model loads, so a bad URI or a missing token fails within
-seconds rather than after a full calibration.
-
-<details>
-<summary>Uploaded artifacts</summary>
-
-| Artifact | Contents |
-| --- | --- |
-| `command.txt` | The full invocation, copy-pasteable, with credentials masked |
-| `version.txt` | The ModelOpt version that ran |
-| `recipe/resolved_recipe.yaml` | The `--recipe` with its `$import`s expanded, so it stands alone |
-| `logs/hf_ptq.log` | The run's Python stdout/stderr, including the traceback if it crashed |
-| `summary/quant_summary.txt` | The per-quantizer summary (unless `--no-verbose`) |
-| `summary/moe.html` | Per-expert calibration token counts, when the run produces them |
-
-</details>
-
-Every command-line argument is also logged as a searchable param, alongside
-`user` / `hostname` / `modelopt_version` / `git_sha` tags. A run that fails is
-still recorded, with status `FAILED` and its log attached.
-
-Other flags:
-
-- `--mlflow_experiment` — defaults to `$USER/hf_ptq/<checkpoint basename>-<recipe name>`,
-  falling back to `--qformat` when no `--recipe` is used.
-- `--mlflow_run_name` — defaults to the UTC start time, `YYYYmmdd-HHMMSS`.
-- `$MLFLOW_TRACKING_URI` enables tracking on its own; `--mlflow` overrides it. A URI taken
-  from the environment is best-effort — if the client is missing or the server is
-  unreachable the run warns and continues untracked, since the variable is often exported
-  for other tooling. An explicit `--mlflow` fails loudly instead.
-
-Authentication uses MLflow's own environment variables (`MLFLOW_TRACKING_TOKEN`, or
-`MLFLOW_TRACKING_USERNAME` / `MLFLOW_TRACKING_PASSWORD`).
-
-The tracking itself lives in `modelopt.torch.utils.mlflow`
-([`MlflowRunLogger`](../../modelopt/torch/utils/mlflow.py)), so other example scripts can
-record runs the same way; `hf_ptq.py` only supplies the params and artifacts specific to PTQ.
-
-> Note: only the main rank uploads, so `--use_fsdp2` runs produce a single run. The log
-> captures Python output; output written directly by native libraries (NCCL, CUDA) goes to
-> the terminal only. On SLURM, keep the job's own `.out` file for those.
-
 ### Megatron-Bridge Example Script
 
 Please refer to [examples/megatron_bridge/README.md](../megatron_bridge/README.md) for example scripts for PTQ / QAD with Megatron-Bridge which is generally more performant than the Hugging Face scripts.
@@ -693,6 +639,61 @@ After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build c
 - Deployable on [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM), [vLLM](https://github.com/vllm-project/vllm) and [SGLang](https://github.com/sgl-project/sglang)
 - More models coming soon!
 
+## Tracking runs with MLflow
+
+Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow <tracking-uri>`, to record a PTQ
+run on an MLflow server so it can be reproduced later from its MLflow entry alone:
+
+```bash
+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>/
+```
+
+The run is opened *before* the model loads, so a bad URI or a missing token fails within
+seconds rather than after a full calibration.
+
+<details>
+<summary>Uploaded artifacts</summary>
+
+| Artifact | Contents |
+| --- | --- |
+| `command.txt` | The full invocation, copy-pasteable, with credentials masked |
+| `version.txt` | The ModelOpt version that ran |
+| `recipe/resolved_recipe.yaml` | The `--recipe` with its `$import`s expanded, so it stands alone |
+| `logs/hf_ptq.log` | The run's Python stdout/stderr, including the traceback if it crashed |
+| `summary/quant_summary.txt` | The per-quantizer summary (unless `--no-verbose`) |
+| `summary/moe.html` | Per-expert calibration token counts, when the run produces them |
+
+</details>
+
+Every command-line argument is also logged as a searchable param, alongside
+`user` / `hostname` / `modelopt_version` / `git_sha` tags. A run that fails is
+still recorded, with status `FAILED` and its log attached.
+
+Other flags:
+
+- `--mlflow_experiment` — defaults to `$USER/hf_ptq/<checkpoint basename>-<recipe name>`,
+  falling back to `--qformat` when no `--recipe` is used.
+- `--mlflow_run_name` — defaults to the UTC start time, `YYYYmmdd-HHMMSS`.
+- `$MLFLOW_TRACKING_URI` enables tracking on its own; `--mlflow` overrides it. A URI taken
+  from the environment is best-effort — if the client is missing or the server is
+  unreachable the run warns and continues untracked, since the variable is often exported
+  for other tooling. An explicit `--mlflow` fails loudly instead.
+
+Authentication uses MLflow's own environment variables (`MLFLOW_TRACKING_TOKEN`, or
+`MLFLOW_TRACKING_USERNAME` / `MLFLOW_TRACKING_PASSWORD`).
+
+The tracking itself lives in `modelopt.torch.utils.mlflow`
+([`MlflowRunLogger`](../../modelopt/torch/utils/mlflow.py)), so other example scripts can
+record runs the same way; `hf_ptq.py` only supplies the params and artifacts specific to PTQ.
+
+> Note: only the main rank uploads, so `--use_fsdp2` runs produce a single run. The log
+> captures Python output; output written directly by native libraries (NCCL, CUDA) goes to
+> the terminal only. On SLURM, keep the job's own `.out` file for those.
+
 ## Resources
 
 - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699)

examples/hf_ptq/example_utils.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
 import copy
 import glob
 import hashlib
@@ -23,13 +24,15 @@
 import shutil
 import warnings
 from collections.abc import Callable, Iterable
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
 from datetime import timedelta
 from pathlib import Path
 from typing import Any
 
 import torch
 import transformers
+import yaml
 from accelerate import infer_auto_device_map, init_empty_weights
 from accelerate.utils import get_max_memory
 from safetensors import safe_open
@@ -43,6 +46,7 @@
     ProcessorMixin,
 )
 
+from modelopt.recipe import load_recipe
 from modelopt.torch.export.model_utils import is_multimodal_model
 
 try:
@@ -51,6 +55,11 @@
     snapshot_download = None
 
 from modelopt.torch.utils import distributed as dist_utils
+from modelopt.torch.utils.mlflow import (
+    MlflowRunLogger,
+    default_experiment_name,
+    validate_tracking_uri,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -1070,3 +1079,119 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str]
     if isinstance(algo.get("layerwise"), dict) and "checkpoint_dir" in algo["layerwise"]:
         algo["layerwise"]["checkpoint_dir"] = resolved
     return quant_cfg, resolved
+
+
+def add_mlflow_args(parser: argparse.ArgumentParser) -> None:
+    """Add the MLflow tracking flags."""
+    parser.add_argument(
+        "--mlflow",
+        default=None,
+        help=(
+            "Track this run on an MLflow server (e.g. https://<your-mlflow-server>/), "
+            "uploading the command, the resolved recipe, the run log and the quantization "
+            "summaries. MLflow's own $MLFLOW_TRACKING_URI enables tracking without this "
+            "flag, which overrides it. A URI taken from the environment is best-effort: if "
+            "it is unusable the run warns and continues untracked."
+        ),
+    )
+    parser.add_argument(
+        "--mlflow_experiment",
+        default=None,
+        help=(
+            "MLflow experiment name. Default: "
+            "$USER/hf_ptq/<checkpoint basename>-<recipe name, or --qformat if no --recipe>."
+        ),
+    )
+    parser.add_argument(
+        "--mlflow_run_name",
+        default=None,
+        help="MLflow run name. Default: the UTC start time as YYYYmmdd-HHMMSS.",
+    )
+
+
+def resolve_mlflow_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
+    """Settle where tracking is configured from, and name the experiment."""
+    # MLflow's own variable enables tracking on its own; --mlflow overrides it. Only the
+    # flag is a deliberate request, so only the flag is fatal when the URI is unusable: the
+    # variable is commonly exported for unrelated tooling and must not fail a quantization.
+    args.mlflow_required = args.mlflow is not None
+    args.mlflow = args.mlflow or os.environ.get("MLFLOW_TRACKING_URI") or None
+    if args.mlflow:
+        try:
+            args.mlflow = validate_tracking_uri(args.mlflow)
+        except ValueError as e:
+            if args.mlflow_required:
+                parser.error(f"--mlflow: {e}")
+            warnings.warn(f"Ignoring MLFLOW_TRACKING_URI, continuing untracked: {e}")
+            args.mlflow = None
+        else:
+            args.mlflow_experiment = args.mlflow_experiment or default_experiment_name(
+                "hf_ptq",
+                args.pyt_ckpt_path,
+                Path(args.recipe).stem if args.recipe else args.qformat,
+            )
+
+
+_MLFLOW_NON_PARAM_ARGS = frozenset(
+    {"dist_state", "mlflow", "mlflow_experiment", "mlflow_required", "mlflow_run_name"}
+)
+
+
+def _mlflow_run_inputs(args: argparse.Namespace) -> tuple[dict, dict]:
+    """Params and start-time artifacts describing this PTQ run."""
+    params = {k: v for k, v in vars(args).items() if k not in _MLFLOW_NON_PARAM_ARGS}
+    # dist_state is an object, so record the one field worth searching on.
+    params["world_size"] = args.dist_state.world_size
+    texts = {}
+    if args.recipe:
+        # The resolved recipe, not the source file: a recipe may be a directory or use
+        # $imports, and only the resolved form is self-contained.
+        resolved = load_recipe(args.recipe).model_dump(mode="json")
+        texts["recipe/resolved_recipe.yaml"] = yaml.safe_dump(resolved, sort_keys=False)
+    return params, texts
+
+
+def _mlflow_logger(args: argparse.Namespace) -> MlflowRunLogger:
+    """Build this run's logger; inert unless --mlflow was given and this is the main rank."""
+    return MlflowRunLogger(
+        args.mlflow,
+        args.mlflow_experiment,
+        run_name=args.mlflow_run_name,
+        enabled=bool(args.mlflow) and args.dist_state.is_main,
+        required=args.mlflow_required,
+    )
+
+
+def mlflow_run(args: argparse.Namespace) -> AbstractContextManager:
+    """Track this invocation for the duration of the block, or do nothing if untracked."""
+    logger = _mlflow_logger(args)
+    if not logger.enabled:
+        # Gathering the inputs re-reads the recipe, so keep it off the untracked path.
+        return nullcontext()
+    params, texts = _mlflow_run_inputs(args)
+    return logger.track(
+        params=params,
+        tags=_mlflow_run_tags(args),
+        texts=texts,
+        files=_mlflow_run_outputs(args),
+    )
+
+
+def _mlflow_run_tags(args: argparse.Namespace) -> dict[str, str]:
+    """Tags shared with the evaluation side, so a PTQ run and the evaluations of the
+    checkpoint it produced can be found together on one tracking server."""
+    return {"model": Path(args.pyt_ckpt_path).name, "checkpoint_path": args.pyt_ckpt_path}
+
+
+def _mlflow_run_outputs(args: argparse.Namespace) -> dict[str, Path]:
+    """Summaries written by post_quantize, keyed by artifact path.
+
+    Uploaded without the leading dot, which is awkward to browse in the MLflow UI. Missing
+    entries are skipped: the MoE table only exists for MoE models, and neither file is
+    written under ``--no-verbose``.
+    """
+    export_path = Path(args.export_path)
+    return {
+        "summary/quant_summary.txt": export_path / ".quant_summary.txt",
+        "summary/moe.html": export_path / ".moe.html",
+    }