diff --git a/examples/hf_ptq/requirements.txt b/examples/hf_ptq/requirements.txt index deb09927544..e4756ce34c7 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 +psutil transformers_stream_generator zstandard diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index 84057e468c9..d3ff4df38b1 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -173,7 +173,16 @@ if [[ $TASKS =~ "quant" ]] || [[ ! -d "$SAVE_PATH" ]] || [[ ! $(ls -A $SAVE_PATH else QUANT_SPEC_ARGS="--qformat=${QFORMAT// /,}" fi - python hf_ptq.py \ + # Opt-in memory/utilization sidecar: wraps the run and writes a CSV trace + + # peak summary under SAVE_PATH. Off by default (no behavior change). + MEM_MON_PREFIX=() + if [[ "${MEM_MONITOR:-0}" == "1" ]]; then + MEM_MON_PREFIX=(python scripts/mem_monitor.py \ + --gpus "${CUDA_VISIBLE_DEVICES-all}" \ + --out "$SAVE_PATH/mem_trace.csv" \ + --summary "$SAVE_PATH/mem_peak.txt" --) + fi + "${MEM_MON_PREFIX[@]}" python hf_ptq.py \ --pyt_ckpt_path=$MODEL_PATH \ --export_path=$SAVE_PATH \ --sparsity_fmt=$SPARSITY_FMT \ diff --git a/examples/hf_ptq/scripts/mem_monitor.py b/examples/hf_ptq/scripts/mem_monitor.py new file mode 100644 index 00000000000..a7989de051d --- /dev/null +++ b/examples/hf_ptq/scripts/mem_monitor.py @@ -0,0 +1,439 @@ +# 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. + +"""Sidecar memory/utilization monitor for HF PTQ runs. + +Samples GPU (memory, utilization, power, temperature) and CPU (total/used/free memory, +utilization) at a fixed interval while +a *separate* workload process (e.g. ``hf_ptq.py``) runs, writes a CSV timeseries +(overwriting any existing file), and prints a peak/mean/min summary on exit. This keeps profiling out of the workload +itself, so it can verify per-run budgets (e.g. the single-GPU layerwise target of +<=80 GB GPU / <=80 GB CPU) without perturbing calibration. + +GPU memory is read at the *device* level (via NVML, falling back to ``nvidia-smi``) +so the monitor observes the workload's usage against the physical device budget +even though it runs as its own process. If no GPU driver is reachable, GPU columns +are left blank and only CPU is reported. Indices are physical NVML/PCI indices; if a +node's CUDA device order differs from NVML order, set ``CUDA_DEVICE_ORDER=PCI_BUS_ID`` +so ``--gpus`` matches the GPUs the workload actually uses. + +Two ways to bind the monitor to a workload: + +* **Wrap mode** (preferred) — pass the workload after ``--``; the monitor launches + it, tracks its process tree (for RSS + CPU%), and exits with its return code:: + + python scripts/mem_monitor.py --gpus 2,3 --out mem.csv --summary peak.txt -- \\ + python hf_ptq.py ... + +* **Standalone** — run in the background and stop it with a signal, ``--duration``, + or ``--pid`` (tracks that PID's tree and exits when it does):: + + python scripts/mem_monitor.py --gpus 2,3 --out mem.csv & MON=$! + python hf_ptq.py ... + kill "$MON" +""" + +import argparse +import contextlib +import csv +import shutil +import signal +import subprocess +import sys +import time +from collections import namedtuple +from pathlib import Path + +import psutil + +MB = 1024**2 + +GpuSample = namedtuple("GpuSample", ["used", "util", "power", "temp"]) +_EMPTY_GPU = GpuSample(None, None, None, None) +CpuStat = namedtuple( + "CpuStat", ["sys_total", "sys_used", "sys_free", "sys_util", "rss", "proc_util"] +) + + +def _to_number(value: str, cast): + """Parse an ``nvidia-smi`` field, returning None for '[N/A]'/unsupported values.""" + try: + return cast(value) + except ValueError: + return None + + +def _resolve_gpu_indices(spec) -> list[int] | None: + """Parse ``--gpus``: ``none``/empty -> [], ``all`` -> None (all devices), else GPU indices. + + ``spec`` may be a single string (CSV like ``"2,3"``) or a list of argparse tokens + (space-separated like ``["2", "3"]``); both are accepted. Non-integer tokens (e.g. the + GPU-UUID / MIG ids that ``CUDA_VISIBLE_DEVICES`` may hold) cannot be mapped to NVML + indices, so GPU monitoring is disabled with a warning rather than crashing the workload. + """ + if isinstance(spec, (list, tuple)): + spec = ",".join(spec) + spec = (spec or "").strip() + if spec.lower() in ("", "none"): + return [] + if spec.lower() == "all": + return None + try: + return [int(x) for x in spec.split(",") if x.strip()] + except ValueError: + print( + f"mem_monitor: --gpus={spec!r} is not integer indices " + "(UUID/MIG ids unsupported); disabling GPU monitoring.", + file=sys.stderr, + ) + return [] + + +class GpuSampler: + """Device-level GPU sampler (memory, utilization, power, temperature) backed by NVML, + falling back to ``nvidia-smi``. + + ``indices`` is a list of physical device indices, ``None`` for all devices, or + an empty list to disable GPU sampling. ``self.indices`` holds the indices that + were actually resolved (empty if no driver is reachable). ``sample()`` returns + ``{index: GpuSample}``. + """ + + def __init__(self, indices: list[int] | None): + self.indices: list[int] = [] + self._backend = None + self._handles: dict[int, object] = {} + self._wanted: set[int] = set() + if indices == []: + return + self._init_nvml(indices) or self._init_smi(indices) + + def _init_nvml(self, indices: list[int] | None) -> bool: + try: + # Optional dependency: pynvml (nvidia-ml-py) may be absent; on any failure + # the sampler falls back to parsing ``nvidia-smi`` in _init_smi. + import pynvml + + pynvml.nvmlInit() + count = pynvml.nvmlDeviceGetCount() + wanted = range(count) if indices is None else [i for i in indices if i < count] + self._handles = {i: pynvml.nvmlDeviceGetHandleByIndex(i) for i in wanted} + self.indices = sorted(self._handles) + self._pynvml = pynvml + self._backend = "nvml" + if indices is not None and (missing := [i for i in indices if i >= count]): + print( + f"mem_monitor: requested GPU indices {missing} exceed device " + f"count {count}; not monitored.", + file=sys.stderr, + ) + return True + except Exception: + return False + + def _init_smi(self, indices: list[int] | None) -> bool: + if shutil.which("nvidia-smi") is None: + return False + try: + available = {i for i, *_ in self._query_smi()} + self.indices = sorted(available if indices is None else available.intersection(indices)) + self._wanted = set(self.indices) + self._backend = "smi" + return True + except Exception: + return False + + @staticmethod + def _query_smi() -> list[tuple]: + out = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,memory.used,utilization.gpu,power.draw,temperature.gpu", + "--format=csv,noheader,nounits", + ], + text=True, + ) + rows = [] + for line in out.strip().splitlines(): + index, used_mb, util, power, temp = (part.strip() for part in line.split(",")) + used = _to_number(used_mb, int) # "[N/A]" on some MIG / unsupported devices + rows.append( + ( + int(index), + used * MB if used is not None else None, + _to_number(util, int), + _to_number(power, float), + _to_number(temp, int), + ) + ) + return rows + + def sample(self) -> dict[int, GpuSample]: + """Return ``{device_index: GpuSample}`` for the resolved indices. + + Each metric is read independently so an unsupported one (e.g. utilization or + power on MIG/vGPU) doesn't drop the others. + """ + if self._backend == "nvml": + result = {} + for i, handle in self._handles.items(): + used = util = power = temp = None + with contextlib.suppress(self._pynvml.NVMLError): + used = self._pynvml.nvmlDeviceGetMemoryInfo(handle).used + with contextlib.suppress(self._pynvml.NVMLError): + util = self._pynvml.nvmlDeviceGetUtilizationRates(handle).gpu + with contextlib.suppress(self._pynvml.NVMLError): + power = self._pynvml.nvmlDeviceGetPowerUsage(handle) / 1000 # mW -> W + with contextlib.suppress(self._pynvml.NVMLError): + temp = self._pynvml.nvmlDeviceGetTemperature( + handle, self._pynvml.NVML_TEMPERATURE_GPU + ) + result[i] = GpuSample(used, util, power, temp) + return result + if self._backend == "smi": + return { + i: GpuSample(used, util, power, temp) + for i, used, util, power, temp in self._query_smi() + if i in self._wanted + } + return {} + + +def _sample_cpu(proc: psutil.Process | None) -> CpuStat: + """System memory (total/used/free) + CPU%, and (if ``proc`` given) tree RSS + process CPU%. + + ``sys_free`` is ``virtual_memory().available`` (allocatable memory incl. reclaimable page + cache), not raw ``.free`` -- the right headroom measure for a budget monitor. + """ + vm = psutil.virtual_memory() + sys_util = psutil.cpu_percent(None) + if proc is None: + return CpuStat(vm.total, vm.used, vm.available, sys_util, None, None) + try: + rss = proc.memory_info().rss + for child in proc.children(recursive=True): + with contextlib.suppress(psutil.Error): + rss += child.memory_info().rss + return CpuStat(vm.total, vm.used, vm.available, sys_util, rss, proc.cpu_percent(None)) + except psutil.Error: + return CpuStat(vm.total, vm.used, vm.available, sys_util, None, None) + + +class _Accumulator: + """Tracks running peak, min, and mean of a metric.""" + + def __init__(self): + self.peak = 0 + self.min = None + self._sum = 0.0 + self._count = 0 + + def add(self, value: float) -> None: + self.peak = max(self.peak, value) + self.min = value if self.min is None else min(self.min, value) + self._sum += value + self._count += 1 + + @property + def mean(self) -> float: + return self._sum / self._count if self._count else 0.0 + + @property + def seen(self) -> bool: + return self._count > 0 + + +class _Metrics: + """Time-aggregated accumulators for every sampled metric.""" + + def __init__(self, gpu_indices: list[int]): + self.gpu = { + i: {k: _Accumulator() for k in ("used", "util", "power", "temp")} for i in gpu_indices + } + self.sys_total = _Accumulator() + self.sys_used = _Accumulator() + self.sys_free = _Accumulator() + self.sys_util = _Accumulator() + self.rss = _Accumulator() + self.proc_util = _Accumulator() + + +def _cell(acc: _Accumulator, value, scale: int = 1, ndigits: int | None = None): + """Record ``value`` into ``acc`` and return its CSV cell ("" when the value is missing).""" + if value is None: + return "" + acc.add(value) + return round(value / scale, ndigits) if ndigits is not None else value + + +def _split_command(argv: list[str]) -> tuple[list[str], list[str] | None]: + """Split ``argv`` on the first ``--`` into (monitor_args, wrapped_command_or_None).""" + if "--" not in argv: + return argv, None + idx = argv.index("--") + return argv[:idx], argv[idx + 1 :] + + +def parse_args(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Sidecar GPU/CPU memory, utilization, power & temperature monitor.", + epilog="Append '-- ' to launch and monitor a workload, exiting with its return code.", + ) + parser.add_argument("--interval", type=float, default=1.0, help="Sampling interval in seconds.") + parser.add_argument( + "--gpus", + nargs="+", + default=["all"], + metavar="GPU", + help="GPUs to monitor: 'all', 'none', or physical indices as CSV ('2,3') or " + "space-separated ('2 3').", + ) + parser.add_argument( + "--pid", + type=int, + default=None, + help="Track this PID's process tree (ignored in wrap mode).", + ) + parser.add_argument("--out", default="mem_trace.csv", help="CSV timeseries output path.") + parser.add_argument( + "--summary", default=None, help="Peak/mean summary path (also printed to stdout)." + ) + parser.add_argument( + "--duration", type=float, default=None, help="Optional max run time in seconds." + ) + return parser.parse_args(argv) + + +def _write_summary(path, duration, metrics: _Metrics): + lines = [f"duration_s: {duration:.1f}"] + for i in sorted(metrics.gpu): + acc = metrics.gpu[i] + if acc["used"].seen: + lines.append(f"peak_gpu{i}_used_mb: {acc['used'].peak / MB:.1f}") + if acc["util"].seen: + lines.append(f"mean_gpu{i}_util_pct: {acc['util'].mean:.1f}") + if acc["power"].seen: + lines.append(f"peak_gpu{i}_power_w: {acc['power'].peak:.1f}") + if acc["temp"].seen: + lines.append(f"peak_gpu{i}_temp_c: {acc['temp'].peak:.0f}") + if metrics.sys_total.seen: + lines.append(f"sys_cpu_total_mb: {metrics.sys_total.peak / MB:.1f}") + if metrics.sys_used.seen: + lines.append(f"peak_sys_cpu_used_mb: {metrics.sys_used.peak / MB:.1f}") + if metrics.sys_free.min is not None: + lines.append(f"min_sys_cpu_free_mb: {metrics.sys_free.min / MB:.1f}") + if metrics.sys_util.seen: + lines.append(f"mean_sys_cpu_util_pct: {metrics.sys_util.mean:.1f}") + if metrics.rss.seen: + lines.append(f"peak_proc_rss_mb: {metrics.rss.peak / MB:.1f}") + lines.append(f"mean_proc_cpu_util_pct: {metrics.proc_util.mean:.1f}") + text = "\n".join(lines) + print(text, flush=True) + if path: + Path(path).parent.mkdir(parents=True, exist_ok=True) + Path(path).write_text(text + "\n") + + +def main() -> None: + monitor_argv, command = _split_command(sys.argv[1:]) + args = parse_args(monitor_argv) + gpu = GpuSampler(_resolve_gpu_indices(args.gpus)) + + child = subprocess.Popen(command) if command else None + target_pid = child.pid if child is not None else args.pid + try: + proc = psutil.Process(target_pid) if target_pid else None + except psutil.NoSuchProcess: + print(f"mem_monitor: --pid {target_pid} is not a running process.", file=sys.stderr) + sys.exit(1) + + metrics = _Metrics(gpu.indices) + + stop = False + + def _request_stop(signum, frame): + nonlocal stop + stop = True + + signal.signal(signal.SIGINT, _request_stop) + signal.signal(signal.SIGTERM, _request_stop) + + # Prime cpu_percent so the first real sample reflects the interval, not 0.0. + psutil.cpu_percent(None) + if proc is not None: + with contextlib.suppress(psutil.Error): + proc.cpu_percent(None) + + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + fieldnames = [ + "elapsed_s", + *(f"gpu{i}_{m}" for i in gpu.indices for m in ("used_mb", "util_pct", "power_w", "temp_c")), + "sys_cpu_total_mb", + "sys_cpu_used_mb", + "sys_cpu_free_mb", + "sys_cpu_util_pct", + "proc_rss_mb", + "proc_cpu_util_pct", + ] + start = time.monotonic() + with open(args.out, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + while not stop: + gpu_sample = gpu.sample() + cpu = _sample_cpu(proc) + elapsed = time.monotonic() - start + + row = {"elapsed_s": round(elapsed, 3)} + for i in gpu.indices: + s = gpu_sample.get(i, _EMPTY_GPU) + acc = metrics.gpu[i] + row[f"gpu{i}_used_mb"] = _cell(acc["used"], s.used, MB, 1) + row[f"gpu{i}_util_pct"] = _cell(acc["util"], s.util) + row[f"gpu{i}_power_w"] = _cell(acc["power"], s.power, 1, 1) + row[f"gpu{i}_temp_c"] = _cell(acc["temp"], s.temp) + row["sys_cpu_total_mb"] = _cell(metrics.sys_total, cpu.sys_total, MB, 1) + row["sys_cpu_used_mb"] = _cell(metrics.sys_used, cpu.sys_used, MB, 1) + row["sys_cpu_free_mb"] = _cell(metrics.sys_free, cpu.sys_free, MB, 1) + row["sys_cpu_util_pct"] = _cell(metrics.sys_util, cpu.sys_util) + row["proc_rss_mb"] = _cell(metrics.rss, cpu.rss, MB, 1) + row["proc_cpu_util_pct"] = _cell(metrics.proc_util, cpu.proc_util) + writer.writerow(row) + f.flush() + + if args.duration is not None and elapsed >= args.duration: + break + if child is not None and child.poll() is not None: + break + if child is None and proc is not None and not proc.is_running(): + break + time.sleep(args.interval) + + if child is not None and child.poll() is None: + child.terminate() + try: + child.wait(timeout=10) + except subprocess.TimeoutExpired: + child.kill() + child.wait() + + _write_summary(args.summary, time.monotonic() - start, metrics) + if child is not None: + rc = child.returncode + sys.exit(rc if rc >= 0 else 128 - rc) # 128+signal when the monitor stopped the child + + +if __name__ == "__main__": + main() diff --git a/tests/examples/hf_ptq/test_mem_monitor.py b/tests/examples/hf_ptq/test_mem_monitor.py new file mode 100644 index 00000000000..2d78ecbe30a --- /dev/null +++ b/tests/examples/hf_ptq/test_mem_monitor.py @@ -0,0 +1,191 @@ +# 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. +"""Unit tests for ``examples/hf_ptq/scripts/mem_monitor.py``. + +The module lives next to the example scripts (not inside the ``modelopt`` package), +so we add ``examples/hf_ptq/scripts/`` to ``sys.path`` before importing it. These +tests are CPU-only: GPU sampling is exercised via the disabled (``none``) path. +""" + +import csv +import os +import subprocess +import sys +from pathlib import Path + +import psutil + +_SCRIPTS_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" / "scripts" +if str(_SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(_SCRIPTS_DIR)) + +import mem_monitor as mm + +_SCRIPT = _SCRIPTS_DIR / "mem_monitor.py" + + +# ---------- helpers --------------------------------------------------------- + + +def test_resolve_gpu_indices(): + assert mm._resolve_gpu_indices("none") == [] + assert mm._resolve_gpu_indices("") == [] + assert mm._resolve_gpu_indices("all") is None + assert mm._resolve_gpu_indices("2,3") == [2, 3] + assert mm._resolve_gpu_indices(" 2 , 3 ") == [2, 3] + # UUID / MIG ids can't map to NVML indices -> disable GPU monitoring, never crash. + assert mm._resolve_gpu_indices("GPU-abcd,MIG-0") == [] + # argparse nargs tokens: space-separated (--gpus 2 3) and single CSV token both work. + assert mm._resolve_gpu_indices(["2", "3"]) == [2, 3] + assert mm._resolve_gpu_indices(["2,3"]) == [2, 3] + assert mm._resolve_gpu_indices(["all"]) is None + + +def test_gpus_accepts_space_separated_and_csv(): + # --gpus 0 1 2 3 (space-separated) no longer crashes argparse; CSV still works. + assert mm._resolve_gpu_indices(mm.parse_args(["--gpus", "0", "1", "2", "3"]).gpus) == [ + 0, + 1, + 2, + 3, + ] + assert mm._resolve_gpu_indices(mm.parse_args(["--gpus", "2,3"]).gpus) == [2, 3] + assert mm._resolve_gpu_indices(mm.parse_args([]).gpus) is None # default 'all' + + +def test_gpu_sampler_disabled(): + sampler = mm.GpuSampler([]) + assert sampler.indices == [] + assert sampler.sample() == {} + + +def test_split_command(): + assert mm._split_command(["--gpus", "none"]) == (["--gpus", "none"], None) + monitor_args, command = mm._split_command(["--out", "x.csv", "--", "python", "-c", "pass"]) + assert monitor_args == ["--out", "x.csv"] + assert command == ["python", "-c", "pass"] + + +def test_sample_cpu_system_only(): + stat = mm._sample_cpu(None) + assert stat.sys_total > 0 + assert stat.sys_used > 0 + assert stat.sys_free >= 0 + assert stat.sys_util >= 0.0 + assert stat.rss is None + assert stat.proc_util is None + + +def test_sample_cpu_process_tree(): + stat = mm._sample_cpu(psutil.Process(os.getpid())) + assert stat.rss is not None and stat.rss > 0 + + +def test_accumulator(): + acc = mm._Accumulator() + assert not acc.seen + for v in (10, 30, 20): + acc.add(v) + assert acc.peak == 30 + assert acc.min == 10 + assert acc.mean == 20.0 + assert acc.seen + + +# ---------- end-to-end (subprocess) ----------------------------------------- + + +def _run(args, **kwargs): + kwargs.setdefault("timeout", 30) + return subprocess.run([sys.executable, str(_SCRIPT), *args], **kwargs) + + +def test_standalone_writes_csv_and_summary(tmp_path): + csv_path = tmp_path / "trace.csv" + summary_path = tmp_path / "peak.txt" + _run( + [ + "--gpus", + "none", + "--interval", + "0.05", + "--duration", + "0.2", + "--out", + str(csv_path), + "--summary", + str(summary_path), + ], + check=True, + ) + + with open(csv_path, newline="") as f: + rows = list(csv.DictReader(f)) + assert rows, "expected at least one sample row" + for col in ( + "elapsed_s", + "sys_cpu_total_mb", + "sys_cpu_used_mb", + "sys_cpu_free_mb", + "sys_cpu_util_pct", + "proc_rss_mb", + "proc_cpu_util_pct", + ): + assert col in rows[0] + + summary = summary_path.read_text() + assert "sys_cpu_total_mb:" in summary + assert "peak_sys_cpu_used_mb:" in summary + assert "min_sys_cpu_free_mb:" in summary + assert "mean_sys_cpu_util_pct:" in summary + + +def test_wrap_mode_propagates_exit_code(tmp_path): + ok = _run( + [ + "--gpus", + "none", + "--interval", + "0.05", + "--out", + str(tmp_path / "a.csv"), + "--", + sys.executable, + "-c", + "import time; time.sleep(0.2)", + ], + ) + assert ok.returncode == 0 + + fail = _run( + [ + "--gpus", + "none", + "--interval", + "0.05", + "--out", + str(tmp_path / "b.csv"), + "--", + sys.executable, + "-c", + "import sys; sys.exit(3)", + ], + ) + assert fail.returncode == 3 + + # Wrap mode tracks the child tree, so proc_rss is populated. + with open(tmp_path / "a.csv", newline="") as f: + rows = list(csv.DictReader(f)) + assert rows and rows[0]["proc_rss_mb"] != ""