Add optional MLflow tracking to hf_ptq.py

CHANGELOG.rst

@@ -6,6 +6,9 @@ Changelog
 
 **New Features**
 
+- Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: it uploads the invocation, the ModelOpt version, the run's log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled.
+- Add ``--mlflow <tracking-uri>`` to ``examples/hf_ptq/hf_ptq.py``, which records a PTQ run through the above and additionally uploads the resolved recipe (``$import``\ s expanded) and the quantization summaries. The run is opened before the model loads so an unreachable server fails in seconds instead of after calibration, and a failed run is still recorded with its log attached. The experiment defaults to ``$USER/hf_ptq/<checkpoint basename>-<recipe name or --qformat>`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``.
+
 **Backward Breaking Changes**
 
 **Deprecations**

examples/hf_ptq/README.md

@@ -293,6 +293,58 @@ 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
+
+Pass `--mlflow <tracking-uri>` to record a PTQ run on an MLflow server, so the run 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`.
+- Passing `--mlflow` with no value uses `$MLFLOW_TRACKING_URI`.
+
+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.

examples/hf_ptq/hf_ptq.py

@@ -19,11 +19,13 @@
 import random
 import time
 import warnings
+from contextlib import AbstractContextManager, nullcontext
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import torch
+import yaml
 from accelerate.hooks import remove_hook_from_module
 from cast_mxfp4_to_nvfp4 import apply_to_model as apply_cast_mxfp4_to_nvfp4
 from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static
@@ -88,6 +90,11 @@
     get_supported_datasets,
 )
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
+from modelopt.torch.utils.mlflow import (
+    MlflowRunLogger,
+    default_experiment_name,
+    validate_tracking_uri,
+)
 from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2
 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
@@ -1622,7 +1629,43 @@ def parse_args() -> argparse.Namespace:
         ),
     )
 
+    parser.add_argument(
+        "--mlflow",
+        nargs="?",
+        const=os.environ.get("MLFLOW_TRACKING_URI", ""),
+        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. Pass the flag without a value to use $MLFLOW_TRACKING_URI."
+        ),
+    )
+    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.",
+    )
+
     args = parser.parse_args()
+    if args.mlflow is not None:
+        try:
+            args.mlflow = validate_tracking_uri(args.mlflow)
+        except ValueError as e:
+            parser.error(f"--mlflow: {e}")
+        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,
+        )
+
     if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0):
         parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].")
 
@@ -1658,6 +1701,70 @@ def parse_args() -> argparse.Namespace:
     return args
 
 
+# Derived state and the tracking settings themselves; everything else argparse parsed is a
+# parameter of the run. Deriving the list means a new flag is tracked without touching this.
+_MLFLOW_NON_PARAM_ARGS = frozenset({"dist_state", "mlflow", "mlflow_experiment", "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,
+    )
+
+
+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",
+    }
+
+
 def main(args: argparse.Namespace):
     if not torch.cuda.is_available():
         raise OSError("GPU is required for inference.")
@@ -1668,31 +1775,17 @@ def main(args: argparse.Namespace):
     setup_distributed_args(args)
 
     try:
-        # launch a memory monitor to read the currently used GPU memory.
-        launch_memory_monitor()
+        # Entered inside the try: opening the run is fatal by design, and skipping
+        # cleanup_distributed would leave the other ranks blocked on the first collective
+        # until the NCCL timeout.
+        with _mlflow_run(args):
+            # launch a memory monitor to read the currently used GPU memory.
+            launch_memory_monitor()
 
-        # Force eager execution for all model types.
-        torch.compiler.set_stance("force_eager")
+            # Force eager execution for all model types.
+            torch.compiler.set_stance("force_eager")
 
-        (
-            full_model,
-            language_model,
-            model_type,
-            calibration_only,
-            processor,
-            tokenizer,
-            default_padding_side,
-            default_pad_token,
-            device,
-        ) = load_model(args)
-
-        if args.sparsity_fmt != "dense":
-            # Sparse
-            sparsity_main(args, full_model, tokenizer, device)
-        else:
-            # Quantize
-            quantize_main(
-                args,
+            (
                 full_model,
                 language_model,
                 model_type,
@@ -1702,7 +1795,25 @@ def main(args: argparse.Namespace):
                 default_padding_side,
                 default_pad_token,
                 device,
-            )
+            ) = load_model(args)
+
+            if args.sparsity_fmt != "dense":
+                # Sparse
+                sparsity_main(args, full_model, tokenizer, device)
+            else:
+                # Quantize
+                quantize_main(
+                    args,
+                    full_model,
+                    language_model,
+                    model_type,
+                    calibration_only,
+                    processor,
+                    tokenizer,
+                    default_padding_side,
+                    default_pad_token,
+                    device,
+                )
     finally:
         cleanup_distributed(args)
 

examples/hf_ptq/requirements.txt

@@ -1,5 +1,6 @@
 compressed-tensors
 fire
 flash-attn>=2.6.0
+mlflow-skinny>=2.9
 transformers_stream_generator
 zstandard

modelopt/torch/utils/logging.py

@@ -34,6 +34,7 @@
 
 __all__ = [
     "DeprecatedError",
+    "TeeStream",
     "atomic_print",
     "capture_io",
     "no_stdout",
@@ -219,5 +220,40 @@ def custom_showwarning(message, category, filename, lineno, file=None, line=None
             warnings.showwarning = original_showwarning
 
 
+class TeeStream:
+    """Mirror a text stream to *sink* while passing writes through to *stream*.
+
+    Scripts that report progress with bare ``print()`` have no log file; wrapping
+    ``sys.stdout``/``sys.stderr`` in this is what produces one. Attribute access falls
+    through to the wrapped stream so ``isatty()`` keeps progress bars behaving. Native
+    (C-level) writes go straight to the real file descriptor and are *not* captured.
+    """
+
+    def __init__(self, stream, sink):
+        """Wrap *stream*, mirroring everything written to it into the open file *sink*."""
+        self._stream = stream
+        self._sink = sink
+
+    def write(self, data: str) -> int:
+        """Write to both the original stream and the sink."""
+        self._stream.write(data)
+        if not self._sink.closed:
+            self._sink.write(data)
+        return len(data)
+
+    def flush(self) -> None:
+        """Flush both the original stream and the sink."""
+        self._stream.flush()
+        if not self._sink.closed:
+            self._sink.flush()
+
+    def __getattr__(self, name):
+        # Guard the wrapped attributes themselves: __getattr__ runs whenever they are absent
+        # (during unpickling, or on a copy), and delegating then would recurse forever.
+        if name in ("_stream", "_sink"):
+            raise AttributeError(name)
+        return getattr(self._stream, name)
+
+
 class DeprecatedError(NotImplementedError):
     """Error for deprecated functions."""