[OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe

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()
 

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,18 +117,24 @@ 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
+        # 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():
             print(
                 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()
 
 
@@ -140,8 +147,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

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()

tests/unit/torch/utils/test_memory_monitor.py

@@ -0,0 +1,135 @@
+# 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 threading
+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, 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"):
+            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 self.devices
+
+    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")
+        if self.on_sample is not None:
+            self.on_sample(handle)
+        return SimpleNamespace(used=2 * 1024**3)
+
+
+@pytest.fixture
+def launched():
+    monitors = []
+    yield monitors
+    for monitor in monitors:
+        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)
+
+
+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()
+    # 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 == ""
+    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