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/modelopt/torch/utils/memory_monitor.py b/modelopt/torch/utils/memory_monitor.py index 94ed1cab850..8001ebf4076 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,9 +117,17 @@ 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(): @@ -126,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() @@ -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 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() 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..eac9aa8d9c2 --- /dev/null +++ b/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