Add MLflow tracking flags to megatron_bridge quantize.py

CHANGELOG.rst

@@ -21,6 +21,7 @@ Changelog
 *Megatron Framework (M-LM / M-Bridge)*
 
 - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ for details.
+- Add ``--mlflow <tracking-uri>`` to ``examples/megatron_bridge/quantize.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too), so a Megatron-Bridge PTQ run records the invocation, every argument as a searchable param, the resolved recipe, the master rank's log and the quantizer summary, and writes ``.experiment.json`` into ``--export_megatron_path``. The experiment defaults to ``$USER/megatron_bridge_quantize/<model basename>-<recipe name or --quant_cfg>`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``.
 
 *Misc*
 

examples/hf_ptq/example_utils.py

@@ -55,11 +55,9 @@
     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,
-)
+from modelopt.torch.utils.mlflow import EXPERIMENT_JSON, MlflowRunLogger, drop_experiment_json
+from modelopt.torch.utils.mlflow import add_mlflow_args as _add_mlflow_args
+from modelopt.torch.utils.mlflow import resolve_mlflow_args as _resolve_mlflow_args
 
 logger = logging.getLogger(__name__)
 
@@ -101,11 +99,8 @@
     "*.tgz",
     "*.zip",
 )
-# Dotted like the other sidecars hf_ptq drops in the export directory, so it is ignored by
-# from_pretrained and does not look like part of the model.
-_EXPERIMENT_JSON = ".experiment.json"
 _HF_PTQ_EXPORT_OWNED_FILES = {
-    _EXPERIMENT_JSON,
+    EXPERIMENT_JSON,
     "config.json",
     "hf_quant_config.json",
     "quant_config.json",
@@ -1307,55 +1302,28 @@ def set_layerwise_export_dir(quant_cfg: dict, export_path: str) -> dict:
 
 def add_mlflow_args(parser: argparse.ArgumentParser) -> None:
     """Add the MLflow tracking flags."""
-    parser.add_argument(
-        "--mlflow",
-        default=None,
-        help=(
+    _add_mlflow_args(
+        parser,
+        "hf_ptq",
+        tracks=(
             "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, and writing .experiment.json into --export_path so the checkpoint "
-            "names the run that produced it. 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>."
+            "names the run that produced it."
         ),
-    )
-    parser.add_argument(
-        "--mlflow_run_name",
-        default=None,
-        help="MLflow run name. Default: the UTC start time as YYYYmmdd-HHMMSS.",
+        variant_help="recipe name, or --qformat if no --recipe",
     )
 
 
 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,
-            )
+    _resolve_mlflow_args(
+        args,
+        parser,
+        tool="hf_ptq",
+        model=args.pyt_ckpt_path,
+        variant=Path(args.recipe).stem if args.recipe else args.qformat,
+    )
 
 
 _MLFLOW_NON_PARAM_ARGS = frozenset(
@@ -1418,58 +1386,18 @@ def mlflow_run(args: argparse.Namespace) -> Iterator[None]:
         try:
             yield
         finally:
-            _log_experiment_json(logger, args, export_path)
-
-
-def _log_experiment_json(
-    logger: MlflowRunLogger, args: argparse.Namespace, export_path: Path
-) -> None:
-    """Record which MLflow run produced this checkpoint, in the checkpoint and on the server.
-
-    The tags point from the run to the checkpoint it wrote; this file is the reverse, so a
-    checkpoint found on disk can be traced back to the run that quantized it without
-    searching the server.
-
-    The artifact goes up for any run that opened, so a failure is traceable from the server
-    side. The local copy is written only once ``export_quantized`` has returned, because the
-    file claims authorship of the checkpoint sitting next to it: ``--export_path`` existing
-    proves nothing, since ``print_quant_summary`` creates it before quantization and the
-    directory may hold a valid checkpoint from an earlier attempt whose weights this run
-    never touched.
-
-    There is nothing to record at all when the run never opened, which a URI taken from the
-    environment reaches by design: it disables tracking from inside the block rather than
-    failing the quantization.
-    """
-    info = logger.run_info
-    if not info:
-        return
-    text = json.dumps(info, indent=2) + "\n"
-    logger.log_text(_EXPERIMENT_JSON.removeprefix("."), text)
-    if not args.checkpoint_exported:
-        return
-    try:
-        (export_path / _EXPERIMENT_JSON).write_text(text)
-    except OSError as e:
-        print(f"[mlflow] WARNING: could not write {export_path / _EXPERIMENT_JSON}: {e}")
+            # Only a completed export may claim the checkpoint next to the pointer:
+            # --export_path existing proves nothing, since print_quant_summary creates it
+            # before quantization and it may hold a checkpoint from an earlier attempt.
+            logger.log_experiment_json(export_path if args.checkpoint_exported else None)
 
 
 def _drop_inherited_experiment_json(args: argparse.Namespace, export_path: Path) -> None:
-    """Remove a pointer an untracked export would otherwise inherit.
-
-    A fresh checkpoint written into a reused ``--export_path`` would keep the previous run's
-    pointer, and one quantized from a tracked source checkpoint could be handed that
-    source's pointer. Either way the file would name a run that did not produce these
-    weights. Only a completed export clears it; a failed run leaves whatever checkpoint was
-    already there, pointer included.
-    """
+    """Drop a pointer an untracked export would otherwise inherit; see
+    :func:`~modelopt.torch.utils.mlflow.drop_experiment_json`."""
     if not args.checkpoint_exported or not args.dist_state.is_main:
         return
-    stale = export_path / _EXPERIMENT_JSON
-    try:
-        stale.unlink(missing_ok=True)
-    except OSError as e:
-        print(f"Warning: could not remove stale {stale}: {e}")
+    drop_experiment_json(export_path)
 
 
 def _mlflow_run_tags(args: argparse.Namespace) -> dict[str, str]:

examples/megatron_bridge/README.md

@@ -124,6 +124,23 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati
 > [!NOTE]
 > HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM covers **Qwen3-VL** and **Qwen3.5-VL**. Other VLMs such as Gemma3-VL are saved in Megatron checkpoint format only.
 
+### Tracking runs with MLflow
+
+Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow <tracking-uri>`, to record a `quantize.py` run on an MLflow server:
+
+```bash
+torchrun --nproc_per_node 2 quantize.py \
+    --hf_model_name_or_path Qwen/Qwen3-8B \
+    --recipe general/ptq/nvfp4_default-kv_fp8 \
+    --tp_size 2 \
+    --export_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \
+    --mlflow https://<your-mlflow-server>/
+```
+
+The run opens *before* the model loads, so a bad URI fails in seconds rather than after a full calibration. Only the master rank uploads: the invocation, every argument as a searchable param, the resolved recipe, that rank's log and the quantizer summary — plus `.experiment.json` written into `--export_megatron_path` once the checkpoint is saved, so a checkpoint on disk names the run that produced it. A failed run is still recorded, with its traceback.
+
+`--mlflow_experiment` defaults to `$USER/megatron_bridge_quantize/<model basename>-<recipe name, or --quant_cfg>`, and `--mlflow_run_name` to the UTC start time. Authentication uses MLflow's own environment variables. See the [`hf_ptq` README](../hf_ptq/README.md#tracking-runs-with-mlflow) for the full artifact list and the `$MLFLOW_TRACKING_URI` semantics.
+
 ## Distillation
 
 This section shows how to distill a student model from a teacher model in the Megatron-Bridge framework.

examples/megatron_bridge/mlflow_utils.py

@@ -0,0 +1,154 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MLflow tracking for ``quantize.py``, mirroring ``examples/hf_ptq``.
+
+Every rank parses and validates the same flags, so a typo in the URI fails identically
+everywhere instead of on one rank while the others wait in a collective. Only the master rank
+opens a run, so the log capture and the uploads happen once.
+
+Nothing here imports Megatron, so the tracking can be exercised without it.
+"""
+
+import argparse
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+import yaml
+
+import modelopt.torch.utils.distributed as dist
+from modelopt.recipe import load_recipe
+from modelopt.torch.utils.mlflow import MlflowRunLogger, drop_experiment_json
+from modelopt.torch.utils.mlflow import add_mlflow_args as _add_mlflow_args
+from modelopt.torch.utils.mlflow import resolve_mlflow_args as _resolve_mlflow_args
+
+TOOL_NAME = "megatron_bridge_quantize"
+
+# The tracking settings describe the destination rather than the quantization, and
+# checkpoint_exported is this script's own bookkeeping.
+_NON_PARAM_ARGS = frozenset(
+    {
+        "checkpoint_exported",
+        "mlflow",
+        "mlflow_experiment",
+        "mlflow_required",
+        "mlflow_run_name",
+    }
+)
+
+
+def add_mlflow_args(parser: argparse.ArgumentParser) -> None:
+    """Add the MLflow tracking flags."""
+    _add_mlflow_args(
+        parser,
+        TOOL_NAME,
+        tracks=(
+            "Track this run on an MLflow server (e.g. https://<your-mlflow-server>/), "
+            "uploading the command, the resolved recipe, the run log and the quantizer "
+            "summary, and writing .experiment.json into --export_megatron_path so the "
+            "checkpoint names the run that produced it."
+        ),
+        variant_help="recipe name, or --quant_cfg if no --recipe",
+    )
+
+
+def resolve_mlflow_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
+    """Settle where tracking is configured from, and name the experiment."""
+    _resolve_mlflow_args(
+        args,
+        parser,
+        tool=TOOL_NAME,
+        model=args.hf_model_name_or_path,
+        # ``or "none"``: neither flag is required by the parser, and the run that reaches
+        # get_quant_config without one fails there rather than while being named.
+        variant=Path(args.recipe).stem if args.recipe else (args.quant_cfg or "none"),
+    )
+
+
+def _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 _NON_PARAM_ARGS}
+    # The parallelism flags say how the run was laid out but not how many GPUs it took:
+    # data parallelism is implicit in the launcher's world size.
+    params["world_size"] = dist.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 _run_tags(args: argparse.Namespace) -> dict[str, str]:
+    """Tags shared with ``hf_ptq`` and the evaluation side, so a PTQ run and whatever is
+    later done with the checkpoint it produced can be found together on one server.
+
+    ``checkpoint_path`` is the checkpoint this run *writes*, because that is what
+    ``export_quantized_megatron_to_hf.py`` (and any QAD run) is later pointed at; the input is
+    kept separately. It is resolved because a relative path is useless as a join key.
+    """
+    return {
+        "model": Path(args.hf_model_name_or_path).name,
+        "checkpoint_path": str(Path(args.export_megatron_path).resolve()),
+        "source_checkpoint_path": args.hf_model_name_or_path,
+    }
+
+
+def _run_outputs(args: argparse.Namespace) -> dict[str, Path]:
+    """Summaries written beside the checkpoint, keyed by artifact path.
+
+    Uploaded without the leading dot, which is awkward to browse in the MLflow UI. A missing
+    entry is skipped: the summary is written by the master rank only once quantization
+    has finished.
+    """
+    return {"summary/quant_summary.txt": Path(args.export_megatron_path) / ".quant_summary.txt"}
+
+
+@contextmanager
+def mlflow_run(args: argparse.Namespace) -> Iterator[None]:
+    """Track this invocation for the duration of the block, and keep the checkpoint's
+    provenance pointer honest whether or not the run is tracked."""
+    logger = MlflowRunLogger(
+        args.mlflow or "",
+        args.mlflow_experiment,
+        run_name=args.mlflow_run_name,
+        enabled=bool(args.mlflow) and dist.is_master(),
+        required=args.mlflow_required,
+    )
+    export_path = Path(args.export_megatron_path)
+    if not logger.enabled:
+        # Gathering the inputs re-reads the recipe, so keep it off the untracked path.
+        try:
+            yield
+        finally:
+            if args.checkpoint_exported and dist.is_master():
+                drop_experiment_json(export_path)
+        return
+    params, texts = _run_inputs(args)
+    with logger.track(
+        params=params,
+        tags=_run_tags(args),
+        texts=texts,
+        files=_run_outputs(args),
+    ):
+        try:
+            yield
+        finally:
+            # Only a completed save may claim the checkpoint the pointer sits next to:
+            # --export_megatron_path exists from print_quant_summary onwards, and may hold a
+            # checkpoint from an earlier attempt whose weights this run never wrote.
+            logger.log_experiment_json(export_path if args.checkpoint_exported else None)

examples/megatron_bridge/quantize.py

@@ -63,6 +63,7 @@
 
 import torch
 from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
+from mlflow_utils import add_mlflow_args, mlflow_run, resolve_mlflow_args
 from transformers import AutoProcessor
 
 import modelopt.torch.quantization as mtq
@@ -222,10 +223,18 @@ def get_args() -> argparse.Namespace:
         help="Skip the post-quantization generation sanity check.",
     )
 
+    add_mlflow_args(parser)
+
     args = parser.parse_args()
+    resolve_mlflow_args(args, parser)
 
     print_args(args)
 
+    # Flipped by main() once the Megatron checkpoint is on disk. The MLflow provenance
+    # pointer is gated on it rather than on --export_megatron_path existing, which proves
+    # nothing: print_quant_summary creates that directory before the save.
+    args.checkpoint_exported = False
+
     return args
 
 
@@ -423,6 +432,7 @@ def forward_loop(_model=None):
         hf_tokenizer_path=args.hf_model_name_or_path,
         hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code},
     )
+    args.checkpoint_exported = True
     if is_vlm:
         print_rank_0(
             f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format. To deploy this "
@@ -462,7 +472,10 @@ def forward_loop(_model=None):
     dist.setup()
     args = get_args()
     try:
-        main(args)
+        # Entered inside the try: opening the run is fatal by design, and the peers of a rank
+        # that exits without dist.abort() stay blocked on the first collective.
+        with mlflow_run(args):
+            main(args)
     except BaseException:
         dist.abort()  # peers may be stuck in a collective this rank will never reach
     finally: