Repository navigation
[OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe - #3133
Conversation
|
/ok to test 4872d09 |
|
/claude review |
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review: 3
Please check the settings in the CodeRabbit UI or the
 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| ) | ||
| if _auto_class(hf_config, trust_remote_code=trust_remote_code) is not AutoModelForCausalLM: | ||
| if ( | ||
| not is_fsdp2_text_only_vlm(hf_config) |
There was a problem hiding this comment.
[CRITICAL Algorithm] Exempting text-only VLMs here lets them past the causal-LM refusal, but the "cannot dequantize" guard six lines above only inspects the top-level config:
if getattr(hf_config, "quantization_config", None) is not None:A composite VLM config nests the quantization settings under text_config, and until this change that didn't matter: every VLM config was refused outright by the _auto_class(...) is not AutoModelForCausalLM check below, so no config with a nested quantization_config could ever reach the loader. Now qwen3_5 / qwen3_5_moe configs skip that check, so a pre-quantized Qwen3.5 checkpoint whose settings live under text_config slips past line 549 and reaches parallel_load_and_prepare_fsdp2, which has no quantization guard of its own (grep for quantization_config in this module: line 549 is the only one). The loader reads the tensors as stored, so packed/compressed weights get loaded as if they were plain BF16 — a silently wrong model, not an error.
The project's own equivalent check already knows this — examples/hf_ptq/example_utils.py:validate_fsdp2_supported refuses pack-quantized / compressed-tensors checkpoints by looking at both:
if any(
getattr(cfg, "quantization_config", None) is not None
for cfg in (config, getattr(config, "text_config", None))
):load_pretrained is public API and is reachable without going through hf_ptq, so the library guard should be at least as wide as the example's. Suggested fix — widen line 549 rather than add a parallel check:
if any(
getattr(cfg, "quantization_config", None) is not None
for cfg in (hf_config, getattr(hf_config, "text_config", None))
):
# It reads the tensors as stored: quantized weights would load as if they were not.
raise NotImplementedError(
f"The {route} load cannot dequantize this {model_type} checkpoint; load it unsharded."
)A test asserting that a qwen3_5 config carrying text_config.quantization_config is still refused would fail without this (and pass with the current code only because the causal-LM check used to catch it).
There was a problem hiding this comment.
Fixed in c51b5d2: the dequantize check reads quantization_config off the text config too, as validate_fsdp2_supported does, and still runs inside the load checks every rank runs together. New test test_an_fsdp_plan_refuses_a_text_only_vlm_quantized_in_its_text_config (qwen3_5 and qwen3_5_moe with text_config.quantization_config set): both cases fail with the old check and pass with the new one.
| # the monitor would never start anywhere. | ||
| try: | ||
| nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) | ||
| nvmlInit() |
There was a problem hiding this comment.
[SUGGESTION] The fix itself is right (probe init → use → finally: nvmlShutdown() keeps pynvml's refcount balanced, and GPUMemoryMonitor.__init__ re-inits for the monitor's own reference), but it's worth noting what it turns on, since this is the first time the monitor actually starts.
Both call sites are unconditional and per-process — examples/hf_ptq/hf_ptq.py:1545 and examples/llm_sparsity/attention_sparsity/hf_sa.py:147 — while _monitor_loop samples every visible GPU and stop() prints one line per GPU. On an 8-rank torchrun job (exactly what this PR series targets) that becomes 8 daemon threads each polling all 8 GPUs every second, and 8 identical ######## blocks at exit. The printed number is already host-wide ("for all processes on the GPU"), so the 7 extra copies carry no information.
A rank gate at the hf_ptq call site would keep the restored monitoring useful without the duplication:
# launch a memory monitor to read the currently used GPU memory.
if args.dist_state.is_main:
launch_memory_monitor()Out of scope for this diff if you'd rather land it with [2/3], which already touches hf_ptq — just flagging it so the behavior change doesn't land unnoticed.
There was a problem hiding this comment.
Not gated here: the gate lands with the hf_ptq change, #2754 ([3/3] of this split), which starts the monitor once per node, on its local rank 0 (args.dist_state.world_size == 1 or local_rank() == 0, tested by test_one_rank_per_node_runs_the_gpu_memory_monitor). It is per node rather than is_main because the monitor reads every GPU on its own node, so rank 0 alone would leave the other nodes of a multi-node run unreported. If this PR merges first, a --use_fsdp2 run on several GPUs starts one monitor per rank until #2754 lands; the duplicates read the same host-wide numbers and change nothing else.
There was a problem hiding this comment.
Claude review — 1 blocking issue
Findings: CRITICAL: 1, IMPORTANT: 0, SUGGESTION: 1
Full scope reviewed: all 6 changed files (3 source + 3 test, 138 added lines). No scoping instructions in the trigger comment.
CRITICAL — the dequantize guard is now narrower than the route it protects
model_load_utils.py:555 — the is_fsdp2_text_only_vlm exemption is correct in itself (build_meta_causal_lm at line 246 already routes these configs to AutoModelForImageTextToText, and _refuse_sharded_load_options still only accepts auto_class in (None, AutoModelForCausalLM)), but it removes the check that was incidentally keeping composite VLM configs away from the loader.
The "cannot dequantize" guard at line 549 reads quantization_config off the top-level config only. Composite VLM configs nest it under text_config — the project's own validate_fsdp2_supported in examples/hf_ptq/example_utils.py checks (config, getattr(config, "text_config", None)) for exactly that reason. Previously no such config could reach line 549's blind spot, because the causal-LM check below refused every VLM. Now qwen3_5 / qwen3_5_moe walk past both, and parallel_load_and_prepare_fsdp2 has no quantization guard of its own (line 549 is the module's only quantization_config check). A pre-quantized Qwen3.5 checkpoint then has its packed weights read as stored, i.e. loaded as if they were plain BF16 — silently wrong, not an error.
load_pretrained is public API reachable without hf_ptq, so the library guard should be at least as wide as the example's. The inline comment has the suggested one-hunk widening of line 549, plus the test that distinguishes it.
What checked out
distributed_forward_loop(forward=...)—forwardhas a real runtime path, the batch handed to it is the device-moved one, and the_disable_use_cache/torch.no_grad/_release_cached_cuda_memory/_reinstall_a2a_forwardswrapping is unchanged. Backward compatible (optional, defaults to the existingmodel(**batch)).Callableis already imported. TheMappingtype check still applies, which is right for a processor'sBatchFeature.- The NVML probe — the diagnosis matches pynvml's behavior (
_nvmlGetFunctionPointerraisesNVMLError_Uninitializedbefore the library is loaded), and thetry/finallykeeps the refcount balanced on both the success and thenvmlDeviceGetHandleByIndexfailure paths;nvmlInitraising itself correctly skips the shutdown.GPUMemoryMonitor.__init__re-inits, so the monitor holds its own reference. - Tests — each change has a test that fails without it. The memory-monitor test's fixture ordering works out (pytest finalizes
launchedbeforemonkeypatch, sostop()runs against the fake NVML, not the real one). No new unit test does anything slow enough to matter for the 20-minlinuxbudget.
SUGGESTION
memory_monitor.py:147 — non-blocking: this is the first commit where the monitor actually starts, and both call sites are unconditional per-process, so an 8-rank run gets 8 threads polling all 8 GPUs and 8 identical exit blocks of host-wide numbers. A dist_state.is_main gate (here or in [2/3], which already touches hf_ptq) would avoid that.
Risk
Low-to-moderate. Two of the three changes are self-contained and well covered; the risk is concentrated in the third, where loosening one refusal exposed a pre-existing narrow check behind it. The failure mode there is silent rather than loud, which is what makes it blocking.
|
/ok to test c51b5d2 |
|
/claude review |
| # the monitor would never start anywhere. | ||
| try: | ||
| nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)) | ||
| nvmlInit() |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The probe fix is right, but it makes stop()'s non-idempotence reachable for the first time.
Until this commit the probe always raised, so launch_memory_monitor always returned None and atexit.register(monitor.stop) on line 157 never ran. Now it does — and stop() (line 112) ends with an unconditional nvmlShutdown() against the single nvmlInit() that __init__ (line 88) issued. Call it twice and the second nvmlShutdown() hits refcount 0 and raises NVMLError_Uninitialized.
That is exactly what this module's own documented usage does (module docstring, lines 29–31):
>>> monitor = launch_memory_monitor(monitor_interval=1.0)
>>> # Run your GPU operations
>>> monitor.stop() # Will print peak memory usage per GPUThe explicit stop() runs, then the atexit hook runs the same bound method again at exit: the peak-memory block prints a second time (with the same numbers, so it reads like a second GPU report) and nvmlShutdown() raises inside the atexit callback, which CPython reports as a traceback on stderr during interpreter shutdown. On a multi-rank torchrun job that is one spurious traceback per rank at the end of an otherwise successful run — noisy in exactly the place people look for real failures.
The new test is itself evidence the hazard is live: tests/unit/torch/utils/test_memory_monitor.py's launched fixture has to call atexit.unregister(monitor.stop) before its own monitor.stop(), i.e. the test works around the double-stop rather than the code preventing it.
The two in-repo call sites (examples/hf_ptq/hf_ptq.py:1545, examples/llm_sparsity/attention_sparsity/hf_sa.py:147) discard the return value and so only stop once — but launch_memory_monitor is public API in modelopt/torch/utils, and the docstring tells users to hold the handle and stop it.
Suggested fix — make stop() idempotent and drop the hook once it has run, so the test's atexit.unregister becomes unnecessary too:
def stop(self):
"""..."""
if self.monitor_thread is None and not self.is_running:
return # Already stopped; NVML's refcount is balanced once only.
atexit.unregister(self.stop)
self.is_running = False
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()
self.monitor_thread = None
nvmlShutdown()A test that fails without it: launch against the fake NVML, call stop() twice, and assert nvml.inits == 0 and that the second call neither raises nor reprints — the current code raises on the second call.
There was a problem hiding this comment.
Fixed in e2ddb75: stop() now does nothing on a monitor that is already stopped, so an explicit stop() followed by the exit hook prints the report once and shuts NVML down once. The test fixture no longer unregisters the hook. New test test_stopping_a_stopped_monitor_does_nothing stops a monitor twice and checks the second call prints nothing and leaves NVML's count at 0; it fails on the previous stop().
| ): | ||
| # It builds through AutoModelForCausalLM, which refuses most VLMs' configs and, for some, | ||
| # builds the language model alone, whose other towers the export would carry unquantized. | ||
| # The text-only VLMs it builds whole instead, through AutoModelForImageTextToText. |
There was a problem hiding this comment.
[SUGGESTION] The exemption and its comment here are clear, but load_pretrained's own docstring — the public API surface for this route — still describes dense="fsdp" as causal-LM-only (line 608):
dense="fsdp": a causal LM FSDP2-sharded over the whole process group, each rank reading its own layers …
After this change the route also accepts the text-only VLM configs is_fsdp2_text_only_vlm names, built whole through AutoModelForImageTextToText. A reader who hits the refusal on line 565 and then reads the docstring to work out which checkpoints qualify gets a rule that no longer matches the code — and the refusal message itself ("builds a causal LM, which this checkpoint is not") now fires for only some VLMs, which is the confusing case to document.
A clause on that bullet would cover it, e.g. "… a causal LM — or one of the text-only VLMs the loader builds whole, vision tower included (is_fsdp2_text_only_vlm) — FSDP2-sharded over the whole process group …". Non-blocking.
There was a problem hiding this comment.
Done in 82a7a1e: the dense="fsdp" bullet now reads "a causal LM, or a text-only VLM the loader builds whole with its vision tower (is_fsdp2_text_only_vlm), FSDP2-sharded over the whole process group".
There was a problem hiding this comment.
Claude review (round 2) — 1 blocking issue
Findings: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 1
Full scope: all 6 changed files (3 source + 3 test, 160 added / 9 deleted). No scoping instructions in the trigger comment, so this is a full review of the diff against shengliangx/hfep-base-6a-ep-loader.
Previous round's CRITICAL is resolved
The dequantize guard now reads text_config as well (model_load_utils.py:549), matching validate_fsdp2_supported's (config, getattr(config, "text_config", None)) tuple in examples/hf_ptq/example_utils.py:102 exactly. test_an_fsdp_plan_refuses_a_text_only_vlm_quantized_in_its_text_config sets config.text_config.quantization_config and asserts the NotImplementedError, and asserts loaders == {} so the refusal happens before any loader is reached — it fails on the pre-fix code. Good fix, and the guard is now wider than the route it protects rather than narrower.
IMPORTANT — stop()'s double-shutdown becomes reachable now that the monitor starts
memory_monitor.py:147 — the probe fix is correct, but it is also the first commit where launch_memory_monitor returns a monitor, so atexit.register(monitor.stop) on line 157 now actually runs. stop() ends in an unconditional nvmlShutdown() against the one nvmlInit() from __init__, so a second call hits refcount 0 and raises NVMLError_Uninitialized — and the module's own documented example (docstring lines 29–31) calls stop() explicitly, which the atexit hook then repeats at exit. Result: the peak-memory block printed twice with identical numbers, and a traceback on stderr from inside the atexit callback — once per rank on a torchrun job, in exactly the place a reader looks for real failures.
The new test is the tell: its launched fixture calls atexit.unregister(monitor.stop) before monitor.stop(), working around the hazard rather than the code preventing it. The two in-repo call sites discard the handle and so stop once, but launch_memory_monitor is public API under modelopt/torch/utils and the docstring tells users to hold it. The inline comment has the idempotent stop() and the test that distinguishes it; the fix also lets the fixture drop its unregister.
What checked out
distributed_forward_loop(forward=...)— real runtime path at line 945–948, inside the unchanged_reinstall_a2a_forwards/_disable_use_cache/torch.no_grad/_release_cached_cuda_memorywrapping, and it receives the device-moved batch, not the raw one. Backward compatible: optional, andforward is Nonekeeps the exact previousmodel(**batch). TheMappingtype check still runs first, which is what you want for a processor'sBatchFeature.Callablewas already imported.- The NVML probe itself — the diagnosis matches pynvml (
_nvmlGetFunctionPointerraisesNVMLError_Uninitializedbefore the library is loaded), and the nesting is right:try/finallybalances the refcount on both the success and thenvmlDeviceGetHandleByIndexfailure paths, whilenvmlInit()raising skips the shutdown it never earned.GPUMemoryMonitor.__init__takes its own reference, so the probe's shutdown does not strand the monitor. - The FSDP2 VLM exemption — consistent end to end:
build_meta_causal_lm(line 246) switchesmodel_factorytoAutoModelForImageTextToTextfor the same predicate, andparallel_load_and_prepare_fsdp2calls it positionally with noauto_class, so these configs really do build whole._sharded_load_kwargshas a single route (_FSDP2), so the exemption cannot leak into a route that would mis-build them. Decoder layers come fromfsdp2_wrap's module-tree walk, not anum_hidden_layersread, so a composite config needs no extra unwrapping._refuse_sharded_load_optionsstill accepts onlyauto_class in (None, AutoModelForCausalLM), and_auto_classis pure, so short-circuiting past it loses no side effect. - Tests — each of the three changes has a test that fails without it, and the FSDP2 pair asserts on
loadersso it distinguishes "routed to the FSDP2 loader" from "routed anywhere". Nothing new spawns processes or trains, so the 20-minlinuxbudget is unaffected.
Still open from round 1 (not re-raised)
The rank-gating SUGGESTION on memory_monitor.py:147 — both call sites are unconditional per-process, so an 8-rank run now gets 8 threads each polling all 8 GPUs and 8 identical host-wide exit blocks. Non-blocking, and [2/3] already touches hf_ptq if you would rather gate it there.
Risk
Low. The three changes are independent and each is well covered; the one blocking item is a two-line robustness fix in the function this PR is repairing, not a problem with the repair.
|
/ok to test 82a7a1e |
|
/claude review |
| dataloader: DataLoader, | ||
| device: torch.device | str | None = None, | ||
| ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, | ||
| forward: Callable[[dict], object] | None = None, |
There was a problem hiding this comment.
[SUGGESTION] The annotation promises forward a dict, but what it actually receives depends on device: with device=None the loop never rebuilds the batch, so the callback gets the loader's own object — a processor's BatchFeature, say — and only the device is not None branch (line 941) hands it a real dict.
Why it matters: the loop's own type check one line above only requires Mapping, which is deliberately wide so a BatchFeature passes. A caller who reads the annotation and writes a callback doing something dict-specific (batch.pop(...), batch | extra, ** into a function expecting a plain dict is fine, but mutation is not) works under device="cuda" and breaks under device=None — the asymmetry is invisible at the call site. Narrowing the annotation to what is guaranteed documents the contract and keeps mypy honest for callers:
| forward: Callable[[dict], object] | None = None, | |
| forward: Callable[[Mapping[str, Any]], object] | None = None, |
(Mapping and Any are both already imported, line 27 / 31.) Non-blocking either way — the alternative is to normalize unconditionally, but that would allocate a dict per batch for callers who pass no device.
There was a problem hiding this comment.
Left as is for now: it is annotation only, with no runtime path that breaks. Its one caller, hf_ptq's expert-parallel image calibration (#2754), always passes the run's device, so its forward always gets the loop's own dict, which it mutates (create_vlm_calibration_forward renames input_ids for encoder-decoder models). I will narrow the annotation to Mapping[str, Any] if this PR is touched again, rather than restart CI on an approved PR for a type hint.
There was a problem hiding this comment.
Claude review (round 3) — approved
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Claude review passed — no blocking issues found. LGTM.
Full scope: all 6 changed files (3 source + 3 test, 184 added / 13 deleted) against shengliangx/hfep-base-6a-ep-loader. No scoping instructions in the trigger comment.
Both prior rounds' blocking items are resolved
Round 1 CRITICAL (dequantize guard too narrow) — _sharded_load_kwargs now reads quantization_config off (hf_config, getattr(hf_config, "text_config", None)) (model_load_utils.py:549), the same tuple validate_fsdp2_supported uses at examples/hf_ptq/example_utils.py:107. The guard is now at least as wide as the route it protects, and test_an_fsdp_plan_refuses_a_text_only_vlm_quantized_in_its_text_config asserts loaders == {}, so it distinguishes "refused before any loader" from "refused somewhere".
Round 2 IMPORTANT (stop()'s double shutdown) — _stopped is set in __init__ and checked first in stop() (memory_monitor.py:88, 124-126), so the atexit hook repeating a caller's explicit stop() neither reprints the peak-memory block nor decrements NVML's refcount past the one nvmlInit() from __init__. The flag is separate from is_running rather than reusing it, which is right — is_running is also False on a monitor built but never started, and such a monitor should still report and shut down. test_stopping_a_stopped_monitor_does_nothing pins both halves (capsys empty, inits == 0), and the launched fixture dropped its atexit.unregister workaround, which is the tell that the hazard is fixed in the code rather than in the test.
What I verified this round
- The NVML probe — nesting is correct: the inner
try/finallybalances the refcount on both the success and thenvmlDeviceGetHandleByIndex/nvmlDeviceGetMemoryInfofailure paths, while a raisingnvmlInit()skips the shutdown it never earned.nvidia-ml-py>=12is a core dependency (pyproject.toml:37), so the new test file'spynvmlimport is safe in the CPU unit job, and the fake NVML means no realnvmlInitruns on a GPU-less runner.stop()setsis_running = Falseand joins beforenvmlShutdown(), so the polling thread cannot call into a shut-down NVML. - The FSDP2 text-only VLM exemption — consistent end to end.
build_meta_causal_lmswitchesmodel_factorytoAutoModelForImageTextToTexton the same predicate (model_load_utils.py:246), and_auto_classindependently returnsAutoModelForImageTextToTextfor a config with atext_config/vision_config(line 1335-1357), so the exemption and the build agree.parallel_load_and_prepare_fsdp2takes noauto_class(line 1196-1207) and_sharded_load_kwargsnever forwards one, so an explicitly passedAutoModelForCausalLMcannot sneak these configs into a text-model-only build — the route builds through the override either way. Decoder layers come fromfsdp2_wrap's module walk (line 1237), not a top-levelnum_hidden_layersread, and the vision tower's tensors fall out generically intonon_layer_sourcesviamodel_param_names, so a composite config needs no special casing._num_experts_from_configalready recurses intotext_config, so theqwen3_5_moevariant is not mis-sized by the auto plan either. distributed_forward_loop(forward=...)— real runtime path at line 945-948, inside the unchanged_reinstall_a2a_forwards/_disable_use_cache/torch.no_grad/_release_cached_cuda_memorywrapping, and it receives the device-moved batch. Backward compatible: appended, optional, andforward is Nonekeeps the exact previousmodel(**batch).- Tests — each of the three changes has a test that fails without it. Nothing new spawns a process, trains, or loads a real checkpoint, so the 20-min
linuxbudget is unaffected.
SUGGESTION (non-blocking)
distributed.py:898 — forward's annotation says dict, but the callback only gets a real dict when device is given; with device=None it gets the loader's own Mapping (a BatchFeature). Mapping[str, Any] matches the loop's own type check.
Still open from round 1 (not re-raised)
The rank-gating SUGGESTION on memory_monitor.py — hf_ptq.py:1545 and hf_sa.py:147 both call launch_memory_monitor() unconditionally per process, so now that the monitor actually starts, an 8-rank run gets 8 threads each polling all 8 GPUs and 8 identical host-wide exit blocks. Non-blocking, and [2/3] already touches hf_ptq if you would rather gate it there.
Risk
Low. Three independent, small changes, each with a test that fails without it; the two issues earlier rounds found are fixed in the code rather than worked around, and the one remaining item is a type-annotation precision nit.
🤖 Generated with Claude Code
e1ab284 to
d9406fe
Compare
82a7a1e to
0cdd82c
Compare
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: the new shutdown test has a thread-start race that can fail otherwise-correct CI runs.
Needs action:
- Remove the first-sample-dependent assertion in
tests/unit/torch/utils/test_memory_monitor.py:94; clear captured output after the first stop and retain the second-stop/no-output and NVML reference-count checks.
| monitor = mm.launch_memory_monitor(monitor_interval=0.01) | ||
|
|
||
| monitor.stop() | ||
| assert "Peak memory usage" in capsys.readouterr().out |
There was a problem hiding this comment.
Bot comment.
[P2] Avoid requiring a completed sample to test shutdown idempotence. Thread.start() waits for the worker to start, not for _monitor_loop to populate peak_memory. The main thread can call stop() before the worker enters while self.is_running, leaving an empty dictionary and only the delimiters in captured output. This assertion then fails despite a correct implementation. Clear the capture after the first stop and rely on the second-stop/no-output and balanced-NVML assertions, or synchronize on a completed sample with a bounded wait.
There was a problem hiding this comment.
Fixed in 216e6ed: the race is real, since peak_memory fills only once the monitor thread has sampled, and stop() right after launch can come first. The first stop is now checked by the report's ######## frame, which stop() always prints, and the test still asserts that the second stop prints nothing and leaves NVML's count at 0. It still fails when stop() is not idempotent, and it passed 30 of 30 repeated runs.
|
/ok to test 216e6ed |
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: the earlier blockers are fixed, but GPUMemoryMonitor.stop() can still race the polling thread while printing its report.
Needs action:
- Fix shutdown ordering in
modelopt/torch/utils/memory_monitor.py: join before readingpeak_memory, and extend the shutdown test with a bounded, two-GPU in-flight sample case (see inline comment).
No action needed:
- ✔️ Four prior concerns addressed, including the nested-quantization guard and repeated NVML shutdown.
- The shutdown-test edit is justified: it removes first-sample timing dependence while retaining second-stop/no-output and NVML reference-count checks.
| """ | ||
| if self._stopped: | ||
| return | ||
| self._stopped = True |
There was a problem hiding this comment.
Bot comment.
[P2] Join the polling thread before printing the report. Setting is_running = False only stops the next outer while iteration; a worker already inside the GPU for loop can still insert another entry into peak_memory. On a host with at least two GPUs, stopping during the first sampling pass can therefore raise RuntimeError: dictionary changed size during iteration while the report iterates peak_memory.items(). That skips both join() and nvmlShutdown(), and the newly set _stopped flag makes the exit hook skip cleanup too.
Move the existing join() immediately after clearing is_running, before reading/reporting peak_memory. Extend the shutdown test using two fake devices and bounded event synchronization around an in-flight sample; the current one-device fake cannot expose this race.
There was a problem hiding this comment.
Fixed in 1876ae9: stop() now joins the polling thread before it reads and prints peak_memory, so a sampling pass in flight finishes first. New test test_the_report_waits_for_a_sampling_pass_in_flight: two fake GPUs; the poller's second sample is held until stop() calls join() (bounded, 5 s at most), and the report must list both GPUs. It failed in 10 of 10 runs with the old ordering, passed in 30 of 30 with the new one, and runs in milliseconds.
|
/ok to test 1876ae9 |
|
/claude review |
| # 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() |
There was a problem hiding this comment.
[SUGGESTION] The reorder is right — the poller mutates peak_memory inside a pass, so reading it before the join could miss the last GPU or iterate a dict that grows under it, and the two-GPU test pins that. Two non-blocking notes on the shape of the wait, both consequences of this being the first commit where the thread actually runs:
-
The report is now behind an unbounded join.
_monitor_looponly checksis_runningat the loop top, aftertime.sleep(self.monitor_interval), sostop()blocks for up to a full interval before printing anything. That latency is pre-existing (the join was unbounded before, just after the print), but the ordering change means a poller wedged in an NVML call — which does happen around an Xid / ECC / device-reset event, exactly when peak memory is interesting — now loses the report entirely instead of printing it and then blocking. Sincelaunch_memory_monitorregistersstopwithatexit, that is a process that hangs at exit with nothing on stdout, once per rank. -
The sleep tail makes the wait longer than it needs to be.
GPUMemoryMonitor's own default ismonitor_interval=10.0, so a directly-constructed monitor'sstop()sits for up to 10s. Both in-repo call sites (hf_ptq.py:1545,hf_sa.py:147) take thelaunch_memory_monitordefault of 1.0s, so this is latent rather than live today.
An Event for the sleep fixes both — stop() returns as soon as the in-flight pass ends, and a bounded join keeps the report when the poller is wedged:
def __init__(self, monitor_interval: float = 10.0):
...
self.monitor_thread = None
self._stopped = False
self._stop_requested = threading.Event()
nvmlInit()
self.device_count = nvmlDeviceGetCount()
def _monitor_loop(self):
while not self._stop_requested.is_set():
for i in range(self.device_count):
...
self._stop_requested.wait(self.monitor_interval)and in stop():
self._stopped = True
self.is_running = False
self._stop_requested.set()
# The thread can be inside a sampling pass, still adding to peak_memory: let it finish
# first, but report what we have rather than hang here if it is wedged in an NVML call.
if self.monitor_thread:
self.monitor_thread.join(timeout=30)Keep is_running as the public flag it already is. The existing two-GPU test still distinguishes the ordering; a timeout case is not worth a test of its own.
There was a problem hiding this comment.
Agreed it would be nicer; not changed in this PR, since neither point is a defect today:
- A poller wedged in an NVML call hangs the exit either way. Before the reorder,
stop()printed and then blocked in the same unboundedjoin(), and the process still hung at exit. What changes is only whether the report printed first. A bounded join would also need to print a snapshot (dict(self.peak_memory)), since a poller that wakes after the timeout would grow the dict during the print again. - The wait after the last pass is at most one
monitor_interval. Both callers uselaunch_memory_monitor's default of 1 s.
I've noted the event-driven sleep, the bounded join and the snapshot as a follow-up for the memory monitor, rather than grow this PR, which Claude has approved.
There was a problem hiding this comment.
Claude review (round 4) — approved
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Claude review passed — no blocking issues found. LGTM.
Full scope: all 6 changed files (3 source + 3 test, 224 added / 15 deleted) against shengliangx/hfep-base-6a-ep-loader. No scoping instructions in the trigger comment. Round 3 approved at 82a7a1e; the checkout is shallow so I could not diff against it directly, and reviewed the current diff whole, concentrating on what moved since — the stop() shutdown ordering and its new test.
meenchen's bot findings are both addressed
The shutdown test's thread-start race — test_stopping_a_stopped_monitor_does_nothing no longer depends on the poller having taken a first sample. It asserts only on the report's frame ("########", printed whether or not a sample landed), clears the captured output, and then keeps the two checks that carry the actual claim: the second stop() prints nothing and nvml.inits == 0. Nothing in it now races the thread's start.
stop() racing the polling thread while printing — memory_monitor.py:127-137 sets is_running = False, joins, and only then reads peak_memory. That is the right order: _monitor_loop writes peak_memory[i] per device inside a pass, so the old order could print after GPU 0's write but before GPU 1's (missing the last GPU on a multi-GPU host), and dict.items() over a dict gaining its first key mid-iteration raises RuntimeError outright. The join still precedes nvmlShutdown(), so the poller can't call into a shut-down NVML either.
test_the_report_waits_for_a_sampling_pass_in_flight pins it and is deterministic rather than timing-based: on_sample holds the thread at handle 1 (main-thread-guarded, so the probe's own handle-0 read doesn't trip it), the test waits on in_flight before stopping, and the join wrapper sets joining only from inside stop() — after is_running = False — so the released thread finishes GPU 1, sees the flag, and exits. It fails on the pre-reorder code. The 5s Event.wait timeouts make it fail fast rather than hang CI.
Re-verified this round
_stoppedidempotency (round 2's IMPORTANT) still holds and is load-bearing for the new test, which stops its monitor explicitly whilelaunch_memory_monitor'satexithook stays registered for the session — the second call returns beforenvmlShutdown()would take the refcount negative against the real pynvml aftermonkeypatchis undone. Keeping_stoppedseparate fromis_runningremains correct:is_runningis alsoFalseon a monitor built but never started, which should still report and shut down.- The NVML probe —
nvmlInit()outside the innertry/finallyso a raising init skips a shutdown it never earned; thefinallybalances the refcount on both the success path and thenvmlDeviceGetHandleByIndex/nvmlDeviceGetMemoryInfofailure paths.GPUMemoryMonitor.__init__takes its own reference, so the probe's shutdown doesn't strand the monitor. - The FSDP2 text-only VLM exemption (round 1's CRITICAL) — unchanged and still consistent end to end. The dequantize guard at
model_load_utils.py:549-556reads(hf_config, getattr(hf_config, "text_config", None)), the same tuplevalidate_fsdp2_supporteduses, so it is wider than the route it protects.is_fsdp2_text_only_vlmcomes from the base branch and the exemption agrees with the build:build_meta_causal_lm:246switchesmodel_factorytoAutoModelForImageTextToTexton the identical predicate, andparallel_load_and_prepare_fsdp2forwards noauto_class, so these configs cannot be built as a text model alone.test_..._refuses_a_text_only_vlm_quantized_in_its_text_configassertsloaders == {}, distinguishing "refused before any loader" from "refused somewhere". distributed_forward_loop(forward=...)— real runtime path atdistributed.py:945-948, inside the unchanged_reinstall_a2a_forwards/_disable_use_cache/torch.no_grad/_release_cached_cuda_memorywrapping, receiving the device-moved batch after theMappingcheck. Appended, optional, andforward is Nonekeeps the exact previousmodel(**batch).- Test cost — nothing new spawns a process, trains, or loads a checkpoint. The memory-monitor tests are
Event-synchronized at a 0.01s interval, so they contribute milliseconds to the 20-minlinuxbudget, not seconds.
SUGGESTION (non-blocking)
memory_monitor.py:128 — the report now sits behind an unbounded join(), and _monitor_loop checks is_running only after a full time.sleep(monitor_interval). The latency is pre-existing (the join was unbounded before too, just after the print), but the ordering means a poller wedged in an NVML call now loses the report rather than printing it first — and since launch_memory_monitor registers stop with atexit, that is a silent hang at exit, once per rank. An Event for the sleep plus join(timeout=...) gets both the race fix and a prompt, bounded stop; the inline comment has the hunk. Both in-repo call sites take the 1.0s default, so this is latent rather than live.
Still open from earlier rounds (not re-raised)
- Rank-gating
launch_memory_monitor(round 1) —hf_ptq.py:1545andhf_sa.py:147call it unconditionally per process, so an 8-rank run gets 8 threads each polling all 8 GPUs and 8 identical host-wide exit blocks. Non-blocking; [2/3] already touches hf_ptq if you would rather gate it there. forward'sCallable[[dict], object]annotation (round 3) — the callback only gets a realdictwhendeviceis given; withdevice=Noneit gets the loader's ownMapping.
Risk
Low. Three independent, small changes, each with a test that fails without it. Every issue the four rounds raised is now fixed in the code rather than worked around in a test, and the remaining items are a robustness refinement and two annotation/ergonomics nits.
🤖 Generated with Claude Code
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: prior correctness concerns are resolved; human sign-off remains for the justified shutdown-test revision.
Needs action:
- Sign off on the shutdown-test revision in
tests/unit/torch/utils/test_memory_monitor.py: removing sample-timing dependence preserves repeated-stop, no-duplicate-output, and balanced-NVML coverage.
No action needed:
- ✔️ Five prior concerns addressed, including the nested-quantization guard, repeated NVML shutdown, and polling/report race.
d9406fe to
ec77bb8
Compare
1876ae9 to
ce19db7
Compare
|
/ok to test ce19db7 |
ec77bb8 to
1546a5e
Compare
ce19db7 to
d2e4abe
Compare
|
/ok to test d2e4abe |
shengliangxu
changed the title
1546a5e to
f1e8c35
Compare
|
/ok to test 4079156 |
d2e4abe to
4079156
Compare
f1e8c35 to
b8c8df5
Compare
|
/ok to test 3be50e3 |
4079156 to
3be50e3
Compare
b8c8df5 to
304c94f
Compare
…iven 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 <shengliangx@nvidia.com>
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 <shengliangx@nvidia.com>
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 <shengliangx@nvidia.com>
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 <shengliangx@nvidia.com>
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 <shengliangx@nvidia.com>
|
/ok to test be5a45b |
3be50e3 to
be5a45b
Compare
Note
[1/3] of #2754's split (2026-10-30), cut to the 500-line budget: [1/3] these library pieces → [2/3] #3139, hf_ptq's previews and whole-model placement → [3/3] #2754, hf_ptq's expert-parallel run. They merge in that order.
Base:
shengliangx/hfep-base-6a-ep-loader, a branch (not a PR) that merges #2779, #3067 and #2740 on main (#2741 [1/3], #3131, has merged). This PR needs only #2740 (distributed_forward_loop); the diff below is this PR's alone.What does this PR do?
Type of change: New feature and bug fix
Two small library changes that hf_ptq's expert-parallel run (#2754) builds on. Each stands on its own and has its own test.
distributed_forward_loop(..., forward=None). The loop calibrates by callingmodel(**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 optionalforward(batch), and still moves each batch to the device and calibrates with the cache disabled. hf_ptq's expert-parallel image calibration uses it ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754).launch_memory_monitorchecks that NVML can report memory before it starts the monitor, but the probe ran beforenvmlInit, so every probe raisedNVMLError_Uninitializedand the monitor never started on any host where nothing else had initialized NVML. The probe now initializes NVML itself and shuts it down afterwards, keeping NVML's reference count balanced. Since the monitor now starts, itsstop(), which also runs at exit, does nothing on a monitor already stopped.load_pretrained's FSDP2 route moved to [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, which now lets them through to the FSDP2 load and refuses one quantized in its text config.Usage
Testing
tests/unit/torch/utils/test_distributed_calibration.py: two batches run through a given forward, on the device and with the cache disabled.tests/unit/torch/utils/test_memory_monitor.py(new), against a fake reference-counted NVML: the monitor starts when NVML can report memory (this fails on the old code), and is skipped with NVML shut down when it cannot.quantization_configtoo, asvalidate_fsdp2_supporteddoes, so a Qwen3.5 checkpoint quantized in its text config is still refused now that these VLMs pass the causal-LM check. The new test fails without it; pre-commit and the three test files pass (transformers 5.18: 172 passed; 5.5.4: 161 passed, 11 skipped).stop()does nothing, since launch_memory_monitor also runs it at exit (the new test fails without it); theload_pretraineddocstring names the text-only VLMs the FSDP2 route takes. Pre-commit and the three test files pass (transformers 5.18: 173 passed; 5.5.4: 162 passed, 11 skipped).hfep-base-6a-ep-loader(its parents rebased onto main 99c3834): no conflicts, own diff unchanged; pre-commit passes against main and the three test files pass on CPU (transformers 5.18: 175 passed).stop()joins the polling thread before it reads and prints the peaks, so a sampling pass in flight finishes first (on several GPUs it could otherwise miss the last GPU or iterate the samples while they grow). The new two-GPU test fails without the reorder. Pre-commit passes; the three test files pass on CPU (transformers 5.5.4: 165 passed, 11 skipped).hfep-base-6a-ep-loader(main 62a5c48, after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 merged). The conflicts were all with [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's newer revision: its new parametrized test sits where this PR's two FSDP2 tests go (all kept), and it rewroteload_pretrained's FSDP2 bullet, which names the text-only VLMs again in the new wording. The rest of the own diff is unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 197 passed; 5.5.4: 185 passed, 12 skipped), and each of the two FSDP2 tests fails with its fix reverted.hfep-base-6a-ep-loader([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 at 47ca6d1). [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 now carries the text-only Qwen3.5 VLM route and the check for a quantization in a VLM's text config, so this PR's three commits for them and their two tests are gone; the rest of the own diff is unchanged. Pre-commit passes against main;test_distributed_calibration.py,test_memory_monitor.pyandtest_model_load_utils.pypass on CPU (transformers 5.18: 207 passed; 5.5.4: 193 passed, 14 skipped).hfep-base-6a-ep-loader: no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 207 passed; 5.5.4: 193 passed, 14 skipped).hfep-base-6a-ep-loader: no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 209 passed; 5.5.4: 194 passed, 15 skipped).hfep-base-6a-ep-loader: no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 219 passed; 5.5.4: 204 passed, 15 skipped).Before your PR is "Ready for review"
forwardis optional; the memory monitor change is a fix.CONTRIBUTING.md: N/AAdditional Information
Split out of #2754 on review (meenchen's bot, "split this PR into stacked [x/3] PRs"). Merge order: this PR, then [2/3] #3139, then #2754 [3/3].