diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..ef087780448 100755 --- a/CHANGELOG.rst +++ b/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 `` 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/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. + **Backward Breaking Changes** **Deprecations** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index a3967565ae0..1c3b47c744c 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -293,6 +293,58 @@ scripts/huggingface_example.sh --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 ` 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 \ + --recipe general/ptq/nvfp4_default-kv_fp8_cast \ + --export_path \ + --mlflow https:/// +``` + +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. + +
+Uploaded artifacts + +| 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 | + +
+ +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/-`, + 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. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6a4bd476984..53c0096c8ad 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/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:///), " + "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/-." + ), + ) + 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) diff --git a/examples/hf_ptq/requirements.txt b/examples/hf_ptq/requirements.txt index deb09927544..cfb87f7da49 100644 --- a/examples/hf_ptq/requirements.txt +++ b/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 diff --git a/modelopt/torch/utils/logging.py b/modelopt/torch/utils/logging.py index d1a9fc1aef9..85d3b9df18f 100644 --- a/modelopt/torch/utils/logging.py +++ b/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.""" diff --git a/modelopt/torch/utils/mlflow.py b/modelopt/torch/utils/mlflow.py new file mode 100644 index 00000000000..fa9b046476f --- /dev/null +++ b/modelopt/torch/utils/mlflow.py @@ -0,0 +1,527 @@ +# 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. + +"""Record a script run on an MLflow tracking server. + +Lets an example script upload its invocation, configuration, log and outputs so the run can +be reproduced from its MLflow entry alone. ``mlflow`` is an optional dependency, imported +only once tracking is actually enabled. +""" + +import contextlib +import getpass +import logging +import os +import re +import shlex +import shutil +import socket +import sys +import tempfile +import time +import traceback +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath +from typing import Any +from urllib.parse import urlparse + +import modelopt +from modelopt.torch.utils.logging import TeeStream + +__all__ = ["MlflowRunLogger", "current_user", "default_experiment_name", "validate_tracking_uri"] + +# MLflow experiment names are stored in a VARCHAR(256) column by the SQL-backed stores. The +# per-component cap stops one pathological component from crowding out the others; the name +# cap is what actually keeps the result storable. +_MAX_COMPONENT_LEN = 100 +_MAX_NAME_LEN = 250 +_UNSAFE_CHARS = re.compile(r"[^A-Za-z0-9._-]+") + +# Anything uploaded or printed passes through _redact first: a tracking URI may carry +# ``user:token@`` and a caller's own flags may carry a secret. +_SECRET_NAME = re.compile(r"token|api[-_]?key|password|passwd|secret|credential", re.IGNORECASE) +_URI_USERINFO = re.compile(r"(?<=://)[^/\s@]+(?=@)") +_MASK = "***" + + +def _stat_key(path: Path) -> tuple[int, int] | None: + """Identity of a file's contents-in-time, or ``None`` when it does not exist.""" + try: + stat = path.stat() + except OSError: + return None + return (stat.st_mtime_ns, stat.st_size) + + +def _redact(value: Any) -> Any: + """Mask credentials embedded in a URI, leaving non-strings untouched.""" + return _URI_USERINFO.sub(_MASK, value) if isinstance(value, str) else value + + +def _redact_argv(argv: list[str]) -> list[str]: + """Mask the value of any ``--*token*`` style option, and credentials in any URI.""" + redacted: list[str] = [] + mask_next = False + for token in argv: + if mask_next: + # Unconditionally, since a secret may itself start with "-"; an option there + # instead would mean the caller passed no value, which argparse rejects anyway. + redacted.append(_MASK) + elif token.startswith("-") and _SECRET_NAME.search(token): + option, sep, _ = token.partition("=") + redacted.append(option + sep + _MASK if sep else option) + else: + redacted.append(_redact(token)) + mask_next = ( + token.startswith("-") and _SECRET_NAME.search(token) is not None and "=" not in token + ) + return redacted + + +def validate_tracking_uri(uri: str) -> str: + """Validate an MLflow tracking URI and return it without a trailing slash. + + Only ``http(s)`` servers are accepted; MLflow's local ``file:`` / ``sqlite:`` backends + are not a useful destination for a shared record of a run. + + Raises: + ValueError: If *uri* is empty, has no host, or is not an http(s) URL. + """ + if not uri: + raise ValueError( + "MLflow tracking URI is empty; pass one explicitly or set MLFLOW_TRACKING_URI." + ) + parsed = urlparse(uri) + if parsed.scheme not in ("http", "https"): + message = f"MLflow tracking URI must be http(s), got {uri!r}." + if not parsed.scheme: + # Only a bare host is plausibly a forgotten scheme; suggesting https://sqlite:///... + # for a URI that already has one would be nonsense. + message += f" Did you mean https://{uri.lstrip('/')}?" + raise ValueError(message) + if not parsed.netloc: + raise ValueError(f"MLflow tracking URI {uri!r} has no host.") + return uri.rstrip("/") + + +def default_experiment_name(tool: str, model: str, variant: str, user: str | None = None) -> str: + """Build an experiment name of the form ``//-``. + + Only the basename of *model* is used, so a local checkpoint directory and an + ``org/name`` Hugging Face id collapse to the same readable name; *variant* is whatever + distinguishes this run of *tool* on *model*, such as a recipe name or a quantization + format. Each component is reduced to ``[A-Za-z0-9._-]`` so the ``/`` separators stay + meaningful, and *user* defaults to the current user. + + Example: + >>> default_experiment_name("hf_ptq", "/models/Qwen3-0.6B/", "nvfp4", user="alice") + 'alice/hf_ptq/Qwen3-0.6B-nvfp4' + """ + owner = user if user is not None else current_user() + name = ( + f"{_sanitize(owner)}/{_sanitize(tool)}/{_sanitize(Path(model).name)}-{_sanitize(variant)}" + ) + return name[:_MAX_NAME_LEN] + + +def current_user() -> str: + """Return the current username, or ``"unknown"`` if the uid has no passwd entry.""" + try: + return getpass.getuser() + except OSError: # container without a passwd entry for the uid + return "unknown" + + +def _sanitize(component: str) -> str: + """Reduce one experiment-name component to ``[A-Za-z0-9._-]``.""" + cleaned = _UNSAFE_CHARS.sub("_", component).strip("._-") + return cleaned[:_MAX_COMPONENT_LEN] or "unknown" + + +def _git_sha() -> str: + """Short commit of the ModelOpt source, or ``"unknown"`` outside a git checkout. + + Read out of ``.git`` rather than by shelling out to ``git``, which keeps the library + free of subprocess use. Handles worktrees, where ``.git`` is a file pointing at the + real git directory and refs live in the main checkout alongside it. + """ + try: + git_path = Path(__file__).resolve().parents[3] / ".git" + if git_path.is_file(): + git_dir = Path(git_path.read_text().split("gitdir:", 1)[1].strip()) + else: + git_dir = git_path + head = (git_dir / "HEAD").read_text().strip() + if not head.startswith("ref: "): + return head[:9] # detached HEAD + ref = head.removeprefix("ref: ") + # A worktree keeps HEAD locally but shares refs with the checkout named by commondir. + bases = [git_dir] + commondir = git_dir / "commondir" + if commondir.is_file(): + bases.append((git_dir / commondir.read_text().strip()).resolve()) + for base in bases: + if (base / ref).is_file(): + return (base / ref).read_text().strip()[:9] + packed = base / "packed-refs" + if packed.is_file(): + for line in packed.read_text().splitlines(): + sha, _, name = line.partition(" ") + if name.strip() == ref: + return sha[:9] + except (OSError, IndexError): + pass + return "unknown" + + +def _command_text() -> str: + """The invocation, as a copy-pasteable line.""" + lines = [shlex.join([sys.executable, *_redact_argv(sys.argv)])] + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size > 1: + lines += [ + "", + f"# Launched under torchrun with WORLD_SIZE={world_size}, " + f"LOCAL_WORLD_SIZE={os.environ.get('LOCAL_WORLD_SIZE', '?')}. The torchrun " + "wrapper is not part of sys.argv and is therefore not shown above.", + ] + return "\n".join(lines) + "\n" + + +class MlflowRunLogger: + """Record one script invocation as an MLflow run. + + :meth:`start` verifies the server and opens the run *before* the expensive work begins, + so a bad URI or a missing token fails in seconds rather than after hours; it also + uploads the invocation and any configuration passed to it, which keeps a crashed run + useful. :meth:`finish` uploads the captured log plus any outputs and closes the run. + Everything is a no-op when ``enabled`` is false, so callers need no branching. + + While the run is open, ``stdout``/``stderr`` are teed to a file that is uploaded as + ``logs/