From e98194dd9c0a7b8a25dcc76dda01415c0a2b04f7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:03:29 +0000 Subject: [PATCH 1/5] Let distributed_forward_loop run each batch through a forward it is given distributed_forward_loop calibrates by calling model(**batch). A batch that needs its own call into the model cannot go through it: a VLM's image-text batch, say, which runs through the VLM's own forward with the keyword arguments that forward accepts. The loop now takes an optional forward(batch), and keeps moving each batch to the device and calibrating with the cache disabled. hf_ptq's expert-parallel image calibration uses it. The new test runs two batches through a given forward and checks the device and the cache. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 15 ++++++++++----- .../torch/utils/test_distributed_calibration.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index e3803ff8644..a8c6c118cc3 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -910,14 +910,16 @@ def distributed_forward_loop( dataloader: DataLoader, device: torch.device | str | None = None, ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, + forward: Callable[[dict], object] | None = None, ) -> Callable[[torch.nn.Module], None]: """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP - all-reduce / all-to-all collectives fire inside that forward). A model whose plan - (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its experts are sharded - over, for the all-to-all forwards that quantizer insertion drops to be re-installed before the - first forward; any other model passes none. ``mtq.quantize`` hands the + all-reduce / all-to-all collectives fire inside that forward), or ``forward(batch)`` when given, + for batches that need their own call into ``model``, such as a VLM's image-text batches. A + model whose plan (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its + experts are sharded over, for the all-to-all forwards that quantizer insertion drops to be + re-installed before the first forward; any other model passes none. ``mtq.quantize`` hands the ``forward_loop`` callback an inner module that this closure ignores, calling the captured top-level ``model`` instead so the full distributed forward (and its collectives) runs. @@ -955,7 +957,10 @@ def calibrate(_inner_model): k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items() } - model(**batch) + if forward is None: + model(**batch) + else: + forward(batch) # Again for what runs next (amax sync, export collectives), past the forwards' pool. _release_cached_cuda_memory() diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py index 2efef9fe86d..cfdf624172a 100644 --- a/tests/unit/torch/utils/test_distributed_calibration.py +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -479,6 +479,21 @@ def test_forward_loop_moves_tensors_to_the_device_it_is_given(): assert [call["logits_to_keep"] for call in model.calls] == [1, 1] +def test_forward_loop_runs_each_batch_through_the_forward_it_is_given(): + """A batch that needs its own call into the model, as a VLM's image-text batch does, goes + through ``forward``, on the device and with the cache disabled like any other.""" + batches = [{"input_ids": torch.zeros(2, 3)}] * 2 + model = _Recorder() + + distributed_forward_loop( + model, batches, device="meta", forward=lambda batch: model(**batch, own_call=True) + )(model) + + assert [call["own_call"] for call in model.calls] == [True, True] + assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2 + assert model.use_cache_seen == [False, False] + + def test_forward_loop_disables_use_cache_only_while_calibrating(): """KV/state caching is off for the calibration forwards (it breaks hybrid Mamba models).""" model = _Recorder() From 83332a22b78285df5c88440616b9a7db51c16f85 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:03:44 +0000 Subject: [PATCH 2/5] Initialize NVML before probing it in launch_memory_monitor launch_memory_monitor checks that NVML can report memory before starting the monitor, so hosts where it cannot (#374) skip monitoring instead of crashing. But the probe ran before nvmlInit, which only GPUMemoryMonitor's constructor calls, so every NVML call in it raised NVMLError_Uninitialized: the monitor printed "Failed to get GPU memory info: Uninitialized" and never started on any host. hf_ptq.py launches it on every run, so no run has reported its per-GPU peak memory since. The probe now initializes NVML for itself and shuts it down afterwards, so the init count stays balanced with the monitor's own. Hosts where memory info is unsupported still skip the monitor. Unit tests run launch_memory_monitor against a fake, reference-counted NVML: the monitor starts when memory info is available (this fails on the old code), and is skipped with NVML left shut down when it is not. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/memory_monitor.py | 9 +- tests/unit/torch/utils/test_memory_monitor.py | 85 +++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 tests/unit/torch/utils/test_memory_monitor.py diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py index 94ed1cab850..0b1fbbb1fd1 100644 --- a/modelopt/torch/utils/memory_monitor.py +++ b/modelopt/torch/utils/memory_monitor.py @@ -140,8 +140,15 @@ def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor | N Returns: GPUMemoryMonitor: The monitor instance that was launched """ + # Probe that NVML can report memory on this host before starting the monitor. NVML must be + # initialized for the probe itself; without it every call raises NVMLError_Uninitialized, and + # the monitor would never start anywhere. try: - nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) + nvmlInit() + try: + nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) + finally: + nvmlShutdown() except Exception as e: print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.") return None diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py new file mode 100644 index 00000000000..8bb3dcab5bf --- /dev/null +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -0,0 +1,85 @@ +# 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 launch_memory_monitor's NVML availability probe, against a fake NVML.""" + +import atexit +from types import SimpleNamespace + +import pytest + +import modelopt.torch.utils.memory_monitor as mm + + +class _FakeNVML: + """Reference-counted like NVML: every call but init raises while uninitialized.""" + + def __init__(self, memory_info_supported: bool = True): + self.inits = 0 + self.memory_info_supported = memory_info_supported + + def install(self, monkeypatch): + for name in ("Init", "Shutdown", "DeviceGetCount", "DeviceGetHandleByIndex"): + monkeypatch.setattr(mm, f"nvml{name}", getattr(self, name)) + monkeypatch.setattr(mm, "nvmlDeviceGetMemoryInfo", self.DeviceGetMemoryInfo) + + def _require_init(self): + if self.inits <= 0: + raise RuntimeError("Uninitialized") + + def Init(self): # noqa: N802 + self.inits += 1 + + def Shutdown(self): # noqa: N802 + self._require_init() + self.inits -= 1 + + def DeviceGetCount(self): # noqa: N802 + self._require_init() + return 1 + + def DeviceGetHandleByIndex(self, index): # noqa: N802 + self._require_init() + return index + + def DeviceGetMemoryInfo(self, handle): # noqa: N802 + self._require_init() + if not self.memory_info_supported: + raise RuntimeError("Not Supported") + return SimpleNamespace(used=2 * 1024**3) + + +@pytest.fixture +def launched(): + monitors = [] + yield monitors + for monitor in monitors: + atexit.unregister(monitor.stop) + monitor.stop() + + +@pytest.mark.parametrize("supported", [True, False], ids=["reports_memory", "cannot"]) +def test_the_monitor_starts_only_where_nvml_can_report_memory(monkeypatch, launched, supported): + """The probe used to run before NVML was initialized, so it always failed and the monitor + never started on any host. Hosts where NVML cannot report memory still skip the monitor, and + leave NVML shut down; elsewhere the probe's own init is balanced and the monitor holds one.""" + nvml = _FakeNVML(memory_info_supported=supported) + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + if monitor is not None: + launched.append(monitor) + + assert (monitor is not None) is supported + assert nvml.inits == (1 if supported else 0) From d4a7f0ea2acfd5afbcb2861a7c59b0362f26b1af Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:31:50 +0000 Subject: [PATCH 3/5] Stop a GPU memory monitor once, however often stop() is called launch_memory_monitor registers the monitor's stop() to run at exit. Now that the monitor starts, a caller that stops it, as the module's own example does, had it stopped again at exit: the peak-memory report printed twice, and the second NVML shutdown raised inside the exit hook. A stopped monitor's stop() now does nothing, and the test fixture no longer unregisters the hook to work around it. The new test stops a monitor twice; it fails without the change. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/memory_monitor.py | 8 +++++++- tests/unit/torch/utils/test_memory_monitor.py | 17 +++++++++++++++-- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py index 0b1fbbb1fd1..918b161f2a8 100644 --- a/modelopt/torch/utils/memory_monitor.py +++ b/modelopt/torch/utils/memory_monitor.py @@ -85,6 +85,7 @@ def __init__(self, monitor_interval: float = 10.0): self.peak_memory = {} # GPU index to peak memory mapping self.is_running = False self.monitor_thread = None + self._stopped = False nvmlInit() self.device_count = nvmlDeviceGetCount() @@ -116,8 +117,13 @@ def stop(self): GPU that was monitored, and properly shuts down the NVML interface. It will wait for the monitoring thread to complete before returning. - The peak memory usage is displayed in GB for each GPU index. + The peak memory usage is displayed in GB for each GPU index. Stopping a stopped monitor does + nothing, as when a caller stops it and the exit hook :func:`launch_memory_monitor` registered + runs too. """ + if self._stopped: + return + self._stopped = True self.is_running = False # Print peak memory usage print("########") diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py index 8bb3dcab5bf..df146afcad5 100644 --- a/tests/unit/torch/utils/test_memory_monitor.py +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -15,7 +15,6 @@ """Unit tests for launch_memory_monitor's NVML availability probe, against a fake NVML.""" -import atexit from types import SimpleNamespace import pytest @@ -66,7 +65,6 @@ def launched(): monitors = [] yield monitors for monitor in monitors: - atexit.unregister(monitor.stop) monitor.stop() @@ -83,3 +81,18 @@ def test_the_monitor_starts_only_where_nvml_can_report_memory(monkeypatch, launc assert (monitor is not None) is supported assert nvml.inits == (1 if supported else 0) + + +def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys): + """launch_memory_monitor registers stop() to run at exit, so a monitor its caller stopped is + stopped again then: the second stop neither reports again nor shuts NVML down a second time.""" + nvml = _FakeNVML() + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + + monitor.stop() + assert "Peak memory usage" in capsys.readouterr().out + monitor.stop() + + assert capsys.readouterr().out == "" + assert nvml.inits == 0 From b6fd97cd08974a7e8c5b8365472860bc743be96d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:56:15 +0000 Subject: [PATCH 4/5] Check the first stop's report without waiting for a GPU sample The idempotence test asserted that the first stop() printed a peak-memory line, which exists only once the monitor thread has sampled a GPU. stop() right after launch can win that race, so the test could fail on a slow runner. It now checks the report's frame, which stop() always prints, and keeps the checks that the second stop prints nothing and leaves NVML's count at 0. It still fails when stop() is not idempotent. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_memory_monitor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py index df146afcad5..136ee8a2b08 100644 --- a/tests/unit/torch/utils/test_memory_monitor.py +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -91,7 +91,8 @@ def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys): monitor = mm.launch_memory_monitor(monitor_interval=0.01) monitor.stop() - assert "Peak memory usage" in capsys.readouterr().out + # The report's frame, which prints whether or not the monitor thread sampled a GPU yet. + assert "########" in capsys.readouterr().out monitor.stop() assert capsys.readouterr().out == "" From be5a45bb4a89ad5050734d2c614bf464176236f3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:14:42 +0000 Subject: [PATCH 5/5] Wait for the memory monitor's sampling pass before reporting its peaks stop() cleared is_running and printed peak_memory before joining the thread, but the thread could be between two GPUs of a sampling pass. The report then missed the last GPU's sample, or, on a host with several GPUs, iterated peak_memory while the thread added to it and raised before the join and the NVML shutdown. stop() now joins the thread first. The new test holds the second of two fake GPUs' samples until stop() waits for the thread, and checks that the report has both GPUs; it fails without the change, in each of 10 runs, and takes milliseconds. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/memory_monitor.py | 5 ++- tests/unit/torch/utils/test_memory_monitor.py | 40 ++++++++++++++++++- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py index 918b161f2a8..8001ebf4076 100644 --- a/modelopt/torch/utils/memory_monitor.py +++ b/modelopt/torch/utils/memory_monitor.py @@ -125,6 +125,9 @@ def stop(self): return self._stopped = True self.is_running = False + # The thread can be inside a sampling pass, still adding to peak_memory: let it finish first. + if self.monitor_thread: + self.monitor_thread.join() # Print peak memory usage print("########") for gpu_idx, peak_mem in self.peak_memory.items(): @@ -132,8 +135,6 @@ def stop(self): f"GPU {gpu_idx}: Peak memory usage = {peak_mem:.2f} GB for all processes on the GPU" ) print("########") - if self.monitor_thread: - self.monitor_thread.join() nvmlShutdown() diff --git a/tests/unit/torch/utils/test_memory_monitor.py b/tests/unit/torch/utils/test_memory_monitor.py index 136ee8a2b08..eac9aa8d9c2 100644 --- a/tests/unit/torch/utils/test_memory_monitor.py +++ b/tests/unit/torch/utils/test_memory_monitor.py @@ -15,6 +15,7 @@ """Unit tests for launch_memory_monitor's NVML availability probe, against a fake NVML.""" +import threading from types import SimpleNamespace import pytest @@ -25,9 +26,11 @@ class _FakeNVML: """Reference-counted like NVML: every call but init raises while uninitialized.""" - def __init__(self, memory_info_supported: bool = True): + def __init__(self, memory_info_supported: bool = True, devices: int = 1, on_sample=None): self.inits = 0 self.memory_info_supported = memory_info_supported + self.devices = devices + self.on_sample = on_sample # called with a device's handle before its memory is read def install(self, monkeypatch): for name in ("Init", "Shutdown", "DeviceGetCount", "DeviceGetHandleByIndex"): @@ -47,7 +50,7 @@ def Shutdown(self): # noqa: N802 def DeviceGetCount(self): # noqa: N802 self._require_init() - return 1 + return self.devices def DeviceGetHandleByIndex(self, index): # noqa: N802 self._require_init() @@ -57,6 +60,8 @@ def DeviceGetMemoryInfo(self, handle): # noqa: N802 self._require_init() if not self.memory_info_supported: raise RuntimeError("Not Supported") + if self.on_sample is not None: + self.on_sample(handle) return SimpleNamespace(used=2 * 1024**3) @@ -97,3 +102,34 @@ def test_stopping_a_stopped_monitor_does_nothing(monkeypatch, capsys): assert capsys.readouterr().out == "" assert nvml.inits == 0 + + +def test_the_report_waits_for_a_sampling_pass_in_flight(monkeypatch, capsys): + """stop() can come while the thread is between two GPUs of a sampling pass: the report waits for + the pass to finish rather than miss the last GPU, or iterate the samples while they grow.""" + in_flight, joining = threading.Event(), threading.Event() + + def hold_the_second_gpu(handle): + if handle == 1 and threading.current_thread() is not threading.main_thread(): + # The first GPU is sampled; hold the second until stop() waits for the thread (5 s max). + in_flight.set() + joining.wait(timeout=5) + + nvml = _FakeNVML(devices=2, on_sample=hold_the_second_gpu) + nvml.install(monkeypatch) + monitor = mm.launch_memory_monitor(monitor_interval=0.01) + join = monitor.monitor_thread.join + + def join_and_release(*args, **kwargs): + joining.set() + return join(*args, **kwargs) + + monkeypatch.setattr(monitor.monitor_thread, "join", join_and_release) + assert in_flight.wait(timeout=5) + + monitor.stop() + + report = capsys.readouterr().out + assert "GPU 0: Peak memory usage" in report + assert "GPU 1: Peak memory usage" in report + assert nvml.inits == 0