Repository navigation
[OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe - #3133
Closed
shengliangxu wants to merge 6 commits into
Closed
shengliangxu wants to merge 6 commits into
shengliangxu wants to merge 6 commits into
Conversation
Collaborator
Author
|
/ok to test 4872d09 |
Collaborator
Author
|
/claude review |
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Per-batch forward callback modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_calibration.py |
distributed_forward_loop invokes the supplied callback for each batch. Without a callback, it retains the existing model(**batch) behavior. The test checks callback calls, batch devices, and disabled model caching. |
GPU Memory Monitor
| Layer / File(s) | Summary |
|---|---|
NVML and monitor lifecycle modelopt/torch/utils/memory_monitor.py, tests/unit/torch/utils/test_memory_monitor.py |
The launcher initializes NVML before probing GPU 0 and shuts it down after the probe. stop() returns when already stopped and joins the sampling thread before reporting peak usage. Tests cover supported and unsupported probes, repeated stops, and sampling across two GPUs. |
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Feature
Possibly related issues
- NVIDIA/Model-Optimizer issue 1699: The callback supports distributed expert-parallel calibration, which is one part of the issue’s broader roadmap objective for distributed model and data-parallel PTQ.
Merge Risk: 
5a987
-
No actionable merge-blocking risk was identified in the reviewed changes.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 45.45% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title names two changes that both appear in the changeset: an optional forward= callback added to distributed_forward_loop, and a fix to the GPU memory monitor's NVML probe (initialize NVML before… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The PR changes only modelopt/torch/utils/distributed.pyandmodelopt/torch/utils/memory_monitor.pyin the reviewed package. The diff adds notorch.load(..., weights_only=False), `numpy.loa…Full details: Title check
-
Explanation
The title names two changes that both appear in the changeset: an optional forward= callback added to distributed_forward_loop, and a fix to the GPU memory monitor's NVML probe (initialize NVML before the memory check and shut it down afterward). The title is bracketed with a ticket ID and a part marker, which is acceptable. It covers the two main changes, though the stop() hardening and the new tests are not mentioned, which is expected for a title.
Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only
modelopt/torch/utils/distributed.pyandmodelopt/torch/utils/memory_monitor.pyin the reviewed package. The diff adds notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosec. It changes nopyproject.tomlorrequirements.txtdependency files. Existingtorch.loadcalls remainweights_only=True.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 45.45% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 11 functions across 4 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
Generate docstrings
-
- Commit to this branch
- Create a new PR
Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
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.
Collaborator
Author
|
/ok to test c51b5d2 |
Collaborator
Author
|
/claude review |
| ): | ||
| # 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. |
Contributor
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.
Collaborator
Author
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".
Contributor
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.
This was referenced Oct 30, 2026
Collaborator
Author
|
/ok to test 82a7a1e |
Collaborator
Author
|
/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, |
Contributor
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:
Suggested change
| 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.
Collaborator
Author
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.
Contributor
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
shengliangxu
force-pushed
the
shengliangx/hfep-base-6a-ep-loader
branch
from
October 30, 2026 00:53
e1ab284 to
d9406fe
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-14a-hf-ptq-lib
branch
from
November 1, 2026 19:56
8569757 to
4c6342d
Compare
Collaborator
Author
|
Sign-off on the shutdown-test revision Signing off on the revised |
shengliangxu
force-pushed
the
shengliangx/hfep-base-6a-ep-loader
branch
from
November 2, 2026 18:33
f49c702 to
7d7548e
Compare
Collaborator
Author
|
/ok to test bb4fbf6 |
shengliangxu
force-pushed
the
shengliangx/hfep-14a-hf-ptq-lib
branch
from
November 2, 2026 18:34
4c6342d to
bb4fbf6
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-base-6a-ep-loader
branch
from
November 2, 2026 20:05
7d7548e to
078e6d3
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-14a-hf-ptq-lib
branch
from
November 2, 2026 20:07
bb4fbf6 to
cee5132
Compare
Collaborator
Author
|
/ok to test cee5132 |
shengliangxu
changed the base branch from
shengliangx/hfep-base-6a-ep-loader
to
main
November 2, 2026 21:26
…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>
shengliangxu
force-pushed
the
shengliangx/hfep-14a-hf-ptq-lib
branch
from
November 2, 2026 21:27
cee5132 to
313800a
Compare
Collaborator
Author
|
/ok to test 313800a |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3133 +/- ##
==========================================
+ Coverage 78.52% 80.16% +1.63%
==========================================
Files 676 676
Lines 77683 77692 +9
==========================================
+ Hits 61002 62279 +1277
+ Misses 16681 15413 -1268
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
|
Without a device, the callback gets the loader's own batch, which can be any mapping rather than a dict. Annotate it as a Mapping, and say that it gets a dict of its own to change only when a device is given. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
/ok to test 5a98712 |
Collaborator
Author
|
Folded into #3139, together with #2754. |
shengliangxu
added a commit
that referenced
this pull request
Nov 3, 2026
…nd apply the model spec when deciding and loading it (#2780) > [!NOTE] > **Folded in #2746 [3/3] on 2026-11-03**, at the user's request: this PR now also applies the model spec when deciding and loading expert-parallel. #2746's commits are here unchanged (this PR's head is #2746's last head, 22b6d23); #2746 keeps their review history, and this PR covers its Jira ticket, OMNIML-6076, too. > > **Re-parented on 2026-10-30** onto #2741 [2/3] (#3136), the expert-parallel loader's core, when #2741 was split. This PR also adopts #2779's newer collective load checks; see the end of Testing. > > **Base:** `main`, since #2741 [2/3] (#3136) merged on 2026-11-03. The diff below is this PR's alone. It merges after #2741 [3/3] (#2741), whose refusals keep `load_pretrained`'s new expert-parallel route from quantizing unsupported experts wrong. <details><summary>Series layout</summary> - `main`, with #3131 #2741 [1/3], #3134 #2746 [1/3], #3135 #2746 [2/3], #2779, #3067, #2740 and #3136 #2741 [2/3] merged - #2741 [3/3] post-load refusals - #2742 dense sharding → #2743 block-FP8 dequantization - **#2780 expert-parallel load routing, with the model spec applied (this PR; #2746 [3/3] folded in)** - `hfep-base-11a-decompress` (#3137) → #2747 on-the-fly decompression - `hfep-base-11b-kimi` (#2746's commits, now in this PR, + #2747) → #3138 → #2748 → #2749 - `hfep-base-14-hf-ptq` (#2746's commits, now in this PR, + #2742 + #2753 + #3133) → #3139 → #2754 </details> ### What does this PR do? Type of change: New feature #2779 routes `load_pretrained` by a `DistributionPlan`, but refuses a plan that shards the routed experts: the expert-parallel loader (#2741) comes later in the series. With it in, this PR adds that route, and applies the model's spec (#3134), through the applier (#3135), when deciding and loading it: - **An expert-parallel plan loads through `load_and_prepare_ep`**, handed the plan itself. Its arguments are checked with the FSDP2 load's, in `_sharded_load_kwargs` inside #2779's collective "load checks" block; `_load_sharded` then runs the load. That load replicates the dense part; FSDP2-sharding it comes with #2742, a sibling of this PR. Its `freeze` and `require_lm_head` options pass through, and so do the Hub options (`token`, `cache_dir`, `local_files_only`), which its config read and `from_pretrained` use. The sharded loads' refusals from #2779 cover it too: a Nemotron-H checkpoint's MTP head, which `load_and_prepare_ep` does not build, is refused rather than exported unquantized. Until #2742 is in, an expert-parallel plan that also shards the dense part is refused before anything loads, naming the ways out: replicate it anyway, or FSDP2-shard the whole model. A plan left to decide shards the dense part when a whole copy of it and a rank's share of the experts would exceed 60% of a GPU. - **The decision refuses to leave an expert-parallel MoE whole on every rank unasked**, naming the expert count the EP degree must divide and every way out: shard its experts (an EP degree is offered only when one divides both the experts and the world), shard the whole model with FSDP2 (an unquantized checkpoint), load it in a single process, or ask for the replication: `MODELOPT_EP_DEGREE=1` with `MODELOPT_FSDP_SIZE=1`, or a replicated plan handed to `load_pretrained`, which skips the decision. A replication both knobs ask for loads, with a note that every rank holds the whole MoE. So #2737's refusals of routed experts the decision can neither count nor shard by now advise `MODELOPT_EP_DEGREE=1` with `MODELOPT_FSDP_SIZE=1` (`ep_size=1` with `dense='replicate'` for a caller passing the arguments): an EP degree of 1 alone now meets the refusal above. - **An MoE the expert-parallel load cannot take loads as a dense model does**, whole on every rank or FSDP2-sharded, when the EP degree is left to the decision or set to 1. The load may refuse it for the model's spec, the transformers release it needs, its checkpoint format or its plan. A rank-0 warning names the reason, and the fix where there is one (an unquantized checkpoint, a later transformers); a replication both knobs ask for gets only the "as asked" note. An EP degree above 1 asked for such a model is still refused by the load. **The model spec in the decision and the load** (folded in from #2746). Without it, an MoE that only its spec makes expert-parallel, such as DeepSeek-V3 before transformers 5.13, is decided as dense. - **Deciding.** Unless the run rules expert parallelism out (a single process, or an EP degree of 1), `resolve_distribution_plan` applies the model's spec to a copy of the config before rank 0 builds the model on meta. - Every rank applies it inside a collective, so a rank that cannot fails every rank. - An expert-parallel plan returns the prepared copy. Any other plan returns the config as read, which the export would otherwise write the spec's plan into. - Rank 0's knobs decide whether the spec applies, and rank 0 sends that decision as it does the plan: a rank whose own environment skipped the spec would load without the expert-parallel plan the others shard by, and fail once the mesh is up. - **Auto and the floors.** An EP degree left to auto resolves to 1 for a model whose spec refuses expert parallelism, or that needs a later transformers than the one installed: its spec's release or the expert-parallel load's 5.7 floor, whichever is later. - The load's floor holds for a model with no spec too (`_spec_refuses_ep`). - Only an explicit EP degree above 1, or a handed expert-parallel plan, meets the refusal: `MODELOPT_FSDP_SIZE=world` alone plans FSDP2 for GPT-OSS, and for Qwen3-MoE on transformers 5.5. - **The left-whole refusal above sees spec plans.** At an EP degree of 1, which skips the spec, it still finds a spec-supplied plan's experts, through the plan the spec would supply (or the one it states `prepare` attaches), without applying anything. - **An MoE whose spec rules expert parallelism out (GPT-OSS, Llama-4), or whose transformers is below the release its expert parallelism needs, loads as a dense model does** when the EP degree is left to auto or set to 1. It loads whole on every rank, or FSDP2-sharded if it doesn't fit one GPU, with the rank-0 warning above carrying the spec's reason or the upgrade (`moe._spec_refusal`). An EP degree above 1 asked for such a model is still refused, with the way around expert parallelism added (`check_distributed_support`). - **Loading.** The expert-parallel route applies the spec as it checks the load's arguments, inside #2779's collective "load checks" block, so a rank that cannot apply it fails every rank instead of leaving the others in the load. - A load handed the config that deciding prepared does not prepare it again, since a model's preparation can patch its modeling code. It still meets its spec's refusal and transformers floor (`check_distributed_support`). - A `config=` passed for the load with a plan handed in is the caller's own, so an expert-parallel load prepares a copy of it, once. - **Why the spec's `prepare` runs before deciding here, and where it goes in #2748.** The decision builds the model on the meta device and finds and counts its routed experts through the plan. For a model whose plan only its spec supplies, the spec must therefore be applied first. This PR applies the whole spec there, `prepare` included, and documents the interim contract on `EPSpec.prepare`: it must leave a whole-model or FSDP2 load as it would be without it, and may run more than once. No spec in this PR sets `prepare`. - #2748 adds the first specs that set it (Kimi-VL, Kimi-K2.5) and replaces the contract. The decision runs only a pure-import step (`EPSpec.make_importable`) and attaches the spec's stated plan (`prepared_ep_plan`) with no `prepare` run. `prepare` then runs only once the plan is decided expert-parallel: in a collective step after the decision, or in the load after its own checks for a plan handed in. - That puts `prepare` where `models/hf.py`'s `prepare_for_loading` hooks already run: on the route that loads the model it prepares, never on one that does not. The hooks are scoped to the load as a context; `prepare` only runs on that route. This closes the series' second phase: an MoE whose expert-parallel plan ships with transformers, such as Qwen3-MoE, or that only its spec makes expert-parallel, such as DeepSeek-V3, now loads expert-parallel through `load_pretrained` and calibrates. ### Usage ```python # torchrun --nproc_per_node 8 quantize.py from modelopt.torch.utils.distributed import distributed_run from modelopt.torch.utils.plugins.model_load_utils import load_pretrained with distributed_run(): model = load_pretrained("Qwen/Qwen3-30B-A3B") # its plan ships with transformers # A model whose plan only its spec supplies loads the same way: # model = load_pretrained("/path/to/DeepSeek-V3-BF16") # ... mtq.quantize(model, ...), export_hf_checkpoint(model, ...) ``` Either model's routed experts are sharded over the 8 ranks. The DeepSeek-V3 path is an unquantized (BF16) copy. The published `deepseek-ai/DeepSeek-V3` checkpoint is FP8, which this PR refuses before anything loads; the expert-parallel load dequantizes FP8 checkpoints from #2743 on. ### Testing - `tests/unit/torch/utils/test_model_load_utils.py` (where #2737 moved the decision's tests) gains the refusal to replicate an expert-parallel MoE, when no degree divides its experts or an explicit EP degree of 1 leaves the dense part to auto, FSDP2 allowed instead, and the replication both knobs ask for loading. Qwen3-MoE shows it, whose plan ships with transformers 5.7. `test_no_ep_degree_dividing_the_experts_leaves_them_whole` goes: the refusal replaces it. #2737's two refusal tests now set `MODELOPT_FSDP_SIZE=1` in their `set-to-1` rows and check the refusals name it. - `tests/unit/torch/utils/test_load_pretrained_distribution.py` (#2779) gains the load's cases: - an expert-parallel plan loads through `load_and_prepare_ep` with the plan, records the plan, and goes through no other loader; - the sharded loads' refusals (placement, model class, MTP head) apply to it. - #2779's test that an expert-parallel plan is refused is gone. - Mutation-checked: routing an expert-parallel plan elsewhere, dropping the replicated-MoE refusal, and applying it to a sharded plan each fail a test. - `tests/gpu/torch/utils/test_load_pretrained_distribution.py` gains a Qwen3-MoE case: left to decide, `load_pretrained` shards its experts, each rank holding its contiguous share and the rest whole, and records the expert-parallel plan. Not run locally; CI's GPU workflow runs it. - Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full `tests/unit` on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite. - Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite. - Review fixes, 2026-10-27: the refusal names no EP degree when the config hides the expert count. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head. - Review fixes, 2026-10-27, round 2: a replication both distribution knobs ask for loads, with a note; the refusal offers an EP degree only when one divides both the experts and the world; the Hub options reach the expert-parallel load. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head. - Rebased on 2026-10-27 onto main 7c3c868, after #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). - Review suggestions, 2026-10-27 evening: the expert-parallel load passes the Hub options to the checkpoint download too. - Re-cut on 2026-10-29 onto #2741 (was on #2743), with three follow-ups: - the test hunks moved from `test_resolve_distribution_plan.py`, which #2737 deleted, into `test_model_load_utils.py`; - #2737's two "Set MODELOPT_EP_DEGREE=1" refusals name `MODELOPT_FSDP_SIZE=1` too; - the `load_pretrained` docstring says the expert-parallel route replicates the dense part. Pre-commit, mypy included, passes against main. On CPU (torch 2.14, transformers 5.18 and 5.5.4), these pass at this head: `test_model_load_utils.py`, `test_load_pretrained_distribution.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `test_distributed.py`. - Review fix, 2026-10-29, from the re-review after the re-cut: on #2741 the expert-parallel load refuses a sharded dense part, so `_route` now refuses an expert-parallel plan that also FSDP2-shards it, before anything loads, naming the ways out. The routing test runs both dense layouts; the refused one fails without the fix. Pre-commit passes against main, and the unit-test files above pass on transformers 5.18 and 5.5.4. - Review suggestion, 2026-10-29: the refusal says that sharding the dense part under expert parallelism is not supported yet (message only). Pre-commit passes against main. - Re-parented on 2026-10-30 onto #2741 [2/3] (#3136); the former head is pinned. - Adapted to #2779's split of the sharded load: the expert-parallel route's checks run in `_sharded_load_kwargs`, inside the collective "load checks" block, and its load in `_load_sharded`. - The Hub options reach the restructured loader's config read and checkpoint step. - Its tests in the deleted `test_load_pretrained_distribution.py` files moved into `test_model_load_utils.py` (unit and GPU), where #2779 moved the load tests. - Review suggestion: the left-whole refusal's FSDP2 way out now says "for an unquantized causal LM", as `_route`'s does (effe11d). Pre-commit, mypy included, passes against main. CPU tests (`test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py`; torch 2.14): | transformers | Passed | Skipped | Failed | |---|---|---|---| | 5.18 | 210 | 0 | 0 | | 5.5.4 | 182 | 28 | 0 | - Review suggestion, 2026-10-30 (e051283): when the expert-parallel load refuses an option it does not take (such as `device_map`), it names its way to a whole model, `MODELOPT_EP_DEGREE=1` with `MODELOPT_FSDP_SIZE=1` or a replicated plan; the FSDP2 load's text is unchanged. A test row checks it and fails without the change. Pre-commit passes against main; `test_model_load_utils.py` and `test_ep_loader.py` pass on transformers 5.18 (161) and 5.5.4 (139, 22 skipped). - Rebased on 2026-10-30 (1bec6da) onto main 99c3834 through #3136's rebase (4c59a6b); no conflicts, this PR's own diff unchanged. Pre-commit passes against main; `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (405) and 5.5.4 (376, 29 skipped). - On GPUs, 2026-10-30: this PR's expert-parallel route, run as the Qwen3-MoE row (`its-own-plan`) of `test_an_expert_parallel_moe_loads_with_its_experts_sharded` at #2746 [3/3]'s ae518d6, passed. That version of the test also compares each rank's logits with the whole model's. #2746's spec-plan row and the two dense `load_pretrained` tests passed beside it, 4 tests in all. NVIDIA GB200, torch 2.11.0a0, transformers 5.9.0, reruns disabled: | GPUs | Passed | Failed | |---|---|---| | 2 | 4 | 0 | | 4 | 4 | 0 | - Rebased on 2026-10-30 (d78aaf8) onto #3136's restack (b726fb3) after #3131 merged to main (62a5c48). The base under #3136 now carries #2779's newer round, so 7 hunks in 3 files conflicted. All are resolved by applying this PR's change to the new text: - `_load_sharded` takes the route beside #2779's new `_has_own_gpu`. - `load_pretrained`'s docstring: the expert-parallel bullet goes before #2779's reworded FSDP2 bullet, and the dropped "an `ep_size > 1` plan raises" sentence leaves the rest of #2779's paragraph as it is. - The tests keep #2779's new by-name, one-rank whole-load and padded-auto cases beside this PR's. The `CHANGELOG.rst` line keeps this PR's wording. - Range-diff otherwise unchanged. Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (467) and 5.5.4 (425, 31 skipped). - Review round, 2026-10-30 (50e64f0): - The dtype refusal names FSDP2's `mp_policy` for the FSDP2 load only (90c914a). - A rank without a GPU of its own is refused before the expert-parallel load too, as before the FSDP2 load. The remedy names the way to a whole model for the route, from one helper shared with the placement refusal (50e64f0). - `test_a_sharded_load_refuses_what_it_cannot_honor` gains a `no_gpu` row and an `mp_policy` assertion. Each fails without its change. - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (469) and 5.5.4 (427, 31 skipped). - Review round, 2026-10-30 (3c353af): - A whole-model load refuses the options only the expert-parallel load takes (`freeze`, `require_lm_head`) before anything loads, naming that load. Before, they reached `from_pretrained` and failed inside transformers. - `test_a_load_refuses_what_it_cannot_load` gains `ep-option-on-a-whole-load`, which fails without the change. - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (470) and 5.5.4 (428, 31 skipped). - Review round, 2026-10-30 (b1f5d8c): - The left-whole refusal names the knobs as the caller set them: `ep_size`/`dense` arguments, which override the environment knobs, or the knobs (b1f5d8c). A new `explicit_degree_one_as_an_argument` row fails without it. - The docstring says the expert-parallel load freezes the model unless `freeze=False`, unlike the other loads (1072e8f). - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (471) and 5.5.4 (428, 32 skipped). - Review round, 2026-10-30 (4a1e020): - The refusal of an MoE left whole on every rank is a `ValueError`, like the decision's other refusals of what its inputs resolve to (d0a63ff). - A load handed an expert-parallel plan and no config reads the checkpoint's config once, for its model type and for the load, as the FSDP2 load already did (4a1e020). The EP-loader test asserts the load gets that read. - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (471) and 5.5.4 (428, 32 skipped). - Review round, 2026-10-30 (4040ea3): - `freeze` or `require_lm_head` passed as `None` asks nothing, as other options passed as `None` do, rather than reading as False and dropping the LM head (83d9ce0). - The expert-parallel route returns the model unfrozen unless `freeze=True`, as the other routes do (69e56a9). - A decided expert-parallel layout says it was decided when it refuses a whole-model placement, as the FSDP2 route already did (4040ea3). - Each change fails a test when reverted. Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (472) and 5.5.4 (428, 33 skipped). - Restack and review round, 2026-10-30 (8f9c8d9), onto #3136 1491b66. The rebase had no conflicts, and this PR's own commits are unchanged; only the first commit's context moved. - A decided expert-parallel layout now takes the load's own `freeze` and `require_lm_head`, rather than refusing them as a whole-model placement (1db4cc3). A new `its-own-options` row fails without the change. - The docstring of `_refuse_a_rank_without_its_gpu` says what the check adds beside #3136's GPU check in the expert-parallel load, which now refuses a run where only some ranks have a GPU (8f9c8d9). - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (476) and 5.5.4 (430, 35 skipped). - Review round, 2026-10-30 (d466147): - A sharded load now refuses an option that only the other sharded load takes by naming that load, rather than with the whole-model way out, which would refuse the option again. `broadcast_chunk_size=None` asks the FSDP2 load to broadcast every layer at once. The expert-parallel load used to drop it silently and now refuses it as the FSDP2 load's option. The two new rows of `test_a_load_refuses_what_it_cannot_load` each fail without the change. - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (478) and 5.5.4 (432, 35 skipped). - Restack and review round, 2026-10-30 (b692242), onto #3136 56353ef on main 059588b, after #2779 reworked the load checks: - With an expert-parallel load to use it, the plan decision no longer takes an EP degree left to auto as 1, so a decided MoE loads expert-parallel. `test_a_load_it_decides_shards_the_experts` replaces #2779's `test_a_load_it_decides_keeps_the_experts_whole`. - This PR's refusals now use #2779's way out, `_whole_instead`, which now takes the route. For a plan handed in, the refusal names a plan with an EP degree of 1; for a decided plan, it also names the knobs. - Placement options on a decided expert-parallel layout are now refused by #2779's `measured` refusal, which this PR extends to that layout. Only options that neither sharded load takes count as placement. - Defaults now follow one rule, #2779's `_ASKS_NOTHING`. On the expert-parallel load, `broadcast_chunk_size=None` is refused as an FSDP2 option, and `cpu_offload=False` asks for nothing. - The refusal of the other sharded load's options now names a way forward (e98619f). - A launch with one GPU per process loads expert-parallel through #2779's GPU predicate (5da80d4). That case fails under the predicate's older reading. - The changelog's backward-breaking entry now says an MoE loads expert-parallel (b692242). - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (513) and 5.5.4 (463, 39 skipped). - Review round, 2026-10-30 (1692141), onto #3136 75bc2d6: - An EP degree left to the plan decision gives way to FSDP2 for the whole model when the dense part, replicated beside a rank's share of the experts, does not fit one GPU, instead of planning a layout the expert-parallel load refuses (1692141). An EP degree the caller asks for still stands and is refused by the load, naming the ways out. `test_what_a_rank_holds_with_the_dense_part_replicated_must_fit` asserts both outcomes, and its decided case fails without the change. - The restack kept #3136's trimmed `load_and_prepare_ep` docstring and comments, plus this PR's sentence on `hub_options`. - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (515) and 5.5.4 (463, 41 skipped). - Restack and review round, 2026-10-30 (093c652), onto #3136 ab06860 on main d464d30: - An EP degree the decision chose now also gives way to none when the expert-parallel load would refuse the model, by its plan (the load's own style checks, now one function, `_ep_plan_refusal`) or by its checkpoint format. Such an MoE is not refused for being left whole either (9cc4589). `test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it` fails without the change in both of its cases. - The refusal of an expert-parallel plan that FSDP2-shards the dense part says that an EP degree left to the decision gives way (381219c). - #2779's texts written before this load now name it: the auto_class refusal, the no-group refusal and the docstring (093c652). - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (519) and 5.5.4 (464, 44 skipped). - Review round, 2026-10-31 (1960174): - An EP degree the decision chose now gives way for every reason the config tells why the expert-parallel load would refuse the model: a model spec ruling it out (Llama-4's, GPT-OSS's), a transformers below the load's floor or the spec's, the checkpoint format, or the plan's styles. The `its-spec` and `its-transformers` cases fail without the change. - Each give-way now says why on rank 0. The CHANGELOG names only the MoEs the expert-parallel load takes as refused when left whole. - Pre-commit passes against main. `test_model_load_utils.py`, `test_ep_loader.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (521) and 5.5.4 (464, 46 skipped). - Restacked on 2026-10-31 (fac2c56) onto #3136 7747c5d, which moved `test_ep_loader.py` into `test_model_load_utils.py`. This PR's one change to that file, the Hub-options checks in `test_the_weights_load_in_the_checkpoints_dtype_through_its_class`, now sits in the moved section. Its own diff is otherwise unchanged. Pre-commit passes against main. `test_model_load_utils.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (521) and 5.5.4 (464, 46 skipped). - Review round, 2026-10-31 (3fc02b8): - The expert-parallel load's transformers floor is one constant, `_EP_MIN_TRANSFORMERS`, that the decision and the load both read (6210606). - The decision now resolves whether the expert-parallel load refuses an MoE whatever its EP degree, so an EP degree of 1 the caller set no longer meets the left-whole refusal for a model that load refuses (3fc02b8). The give-way test runs each refusal with the degree on auto and with `MODELOPT_EP_DEGREE=1`, and the `set-to-1` cases fail without the change. - Pre-commit passes against main. `test_model_load_utils.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (525) and 5.5.4 (464, 50 skipped). - Review round, 2026-10-31 (3c40b57): - `load_and_prepare_ep` refuses what the plan decision gives way for (`_ep_load_refusal`), after its own floor and format checks. An EP degree asked for, or a plan handed in, no longer loads a model whose spec rules expert parallelism out. The load test's new GPT-OSS case fails without the change. - Pre-commit passes against main. `test_model_load_utils.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (526) and 5.5.4 (465, 50 skipped). - Restacked on 2026-10-31 (c2d18e6) onto #3136 930bb0f, on the new main e352cac and #2779 c78d2e5, with Claude's 04:08Z review round folded in: - #2779's texts meet the expert-parallel load. Its new `init_quantized_weights` refusal names the expert-parallel route's own way to a whole model, and a new test covers both a knob and a handed plan; the plain FSDP2 wording fails it. The loading guide #2779 added now describes the expert-parallel route, which this PR's docstring paragraphs moved into (5455a0c). The decided EP degree is no longer forced to 1 in `load_pretrained`. - `_ep_load_refusal` returns the error itself, so a floor refusal is a `RuntimeError` whichever floor it trips (eda4c88). A spec's floor above the load's own is tested (608b0b8). Both cases fail without their change. - Pre-commit passes against main. `test_model_load_utils.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (536) and 5.5.4 (471, 54 skipped). - Review round, 2026-10-31 (f7b11fe): the dense-sharding give-way test also leaves the EP degree to the decision, where a caller asking for FSDP2 makes the decided degree give way; it fails if only rank 0's GPU memory may. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (537) and 5.5.4 (471, 55 skipped). - Restacked on 2026-11-01 (36863bc) onto #3136 4dd38a3, on main cf38ec3 and #2779 c42cb61. #2779's new whole-model branch and its two-rank test come in unchanged. The CHANGELOG entry takes #2779's rewritten sentence, with this PR's MoE clause in place of the one saying an MoE loads without expert parallelism. Its own diff is otherwise unchanged. Pre-commit passes against main. `test_model_load_utils.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (538) and 5.5.4 (472, 55 skipped). - Restacked on 2026-11-01 (8e19843) onto #3136 35f1ac0, on main 51645db, without conflicts; its own diff and tree are unchanged. Pre-commit passes against main. - Design change, 2026-11-01 (e20bd19): an MoE the expert-parallel load cannot take now loads as a dense model does, with a rank-0 warning, both when the EP degree is left to the decision and when it is set to 1 (before, only the first was noted, on rank 0, and the second loaded silently). The give-way test asserts the warning for each reason the load refuses, and fails without it. Pre-commit passes against main. `test_model_load_utils.py`, `test_moe_a2a.py` and `tests/unit/torch/models` pass on transformers 5.18 (538) and 5.5.4 (472, 55 skipped). - Review round, 2026-11-01 (edf3b61): the dense-load warning skips a replication both knobs ask for, which keeps its "as asked" note. The give-way test asserts no warning there, and fails without the gate. The guide promises a fix only where the load names one. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (543) and 5.5.4 (472, 60 skipped). - Restacked on 2026-11-02 (301f97d) on #3136 39f55f2, on main 2d4e8d0, where #2779 and #3135 merged. There were no conflicts, and the PR's own diff is unchanged. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped). - Restacked on 2026-11-02 (09322c1) on #3136 c0041d0, on main 4f08a12, where #3067 merged. There were no conflicts, and the PR's own diff is unchanged. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped). - Restacked on 2026-11-02 (7653ddf) on #3136 d87f5bd, which now targets main 653ef61, where #2740 merged. There were no conflicts, and the PR's own diff is unchanged. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped). - Rebased on 2026-11-03 (e2d81c9) onto main a6cc167, where #3136 merged; this PR's base is now `main`. There were no conflicts, and the PR's own diff is unchanged. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (572) and 5.5.4 (486, 75 skipped). - **From #2746, folded in on 2026-11-03.** #2746's Testing section keeps its per-round record. Its tests and checks: - `tests/unit/torch/utils/test_model_load_utils.py` (the decision's and the loads' tests) gains the spec's cases: - the spec applied before the decision, DeepSeek-V3's plan coming only from its spec, and the returned config carrying it for an expert-parallel plan only; - no spec for a run that rules expert parallelism out, by argument or by `MODELOPT_EP_DEGREE=1`; - auto picks no EP degree for GPT-OSS, nor below a spec's floor or the load's (DeepSeek-V3 on transformers 5.6.2); an explicit `ep_size=2` meets the refusal; - on two real gloo ranks, every rank applying the spec, or not, as rank 0 decided, whatever its own knob says; - on two real gloo ranks, a rank that cannot apply the spec fails every rank, when deciding and when loading a handed expert-parallel plan; - the left-whole refusal for spec plans at an EP degree of 1; GPT-OSS, and DeepSeek-V3 below the floor, loading as a dense model with the warning, at an EP degree left to auto or set to 1; - a decided expert-parallel load gets the spec-prepared config, prepared once, not at all when handed, and on a copy when passed as `config=` with a handed plan; GPT-OSS is refused with its config handed in, without preparing it. - Mutation-checked (earlier rounds): skipping the spec or applying it without expert parallelism, ignoring the EP knob, letting each rank decide the spec for itself, preparing a handed config twice or not at all, ignoring the spec's plan at an EP degree of 1, letting auto pick an EP degree for a model whose spec refuses expert parallelism or below a floor, and offering an EP degree below the floor each fail a test. - Checked locally, at the stack's tip: DeepSeek-V3, Nemotron-H and Nemotron-H-Omni configs prepared by their specs and swapped to the `torch_a2a` styles are written by the expert-parallel export with no `base_model_ep_plan`, as their sources have none. There the export writer itself puts back the source `config.json`'s values (#2753). The loaded model's own config holds no spec plan only since b51123a (see #2746's Testing). - On GPUs, 2026-10-30, at ae518d6: `test_an_expert_parallel_moe_loads_with_its_experts_sharded` in `tests/gpu/torch/utils/test_model_load_utils.py` runs Qwen3-MoE (`its-own-plan`) and DeepSeek-V3 saved with no plan of its own (`its-spec-plan`) through the real `load_pretrained`, with the plan left to decide. It checks each rank's expert slices and its own batch's logits against the whole model. It and the two dense `load_pretrained` tests passed on 2 NVIDIA GB200 GPUs (torch 2.11.0a0, transformers 5.9.0, reruns disabled). - Folded #2746 in on 2026-11-03 (22b6d23): this PR's branch fast-forwarded to #2746's head, so no commit changed. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (594) and 5.5.4 (495, 88 skipped). - Review fix, 2026-11-03 (c9f3a2d): the refusal of a sharded load inside `init_quantized_weights` gives each route its own reason. The expert-parallel load calls `from_pretrained` with its own config and device mesh, and the FSDP2 load's text is unchanged; the guide and a test docstring follow. Pre-commit passes against main. The same CPU tests pass on transformers 5.18 (594) and 5.5.4 (495, 88 skipped). - Review fix, 2026-11-03 (275ce95): the Backward Breaking CHANGELOG entry on how `load_pretrained` distributes a model is cut to two sentences, what changed and how to keep loading it whole on every rank. Pre-commit passes against main. - Review fix, 2026-11-03 (0476abb): the guide's paragraph on calls that place the model names the ways out the refusal gives, by route. For an MoE whose experts the decision shards, that is `MODELOPT_EP_DEGREE=1` with `MODELOPT_FSDP_SIZE=1`, or a plan with an EP degree of 1. Pre-commit passes against main. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Nothing released changes. Against #2737 and #2779: an expert-parallel plan now loads, and a decision to replicate an expert-parallel MoE on every rank now raises unless both knobs ask for it. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ #2779's 0.49.0 Misc entry now names expert sharding. - Did you get Claude approval on this PR?: This PR was approved on 2026-10-30 (effe11d), and #2746 on 2026-11-01 (04a197e), each with its suggestions taken. Re-review requested for the folded PR. ### Additional Information Part of expert-parallel PTQ for Hugging Face MoE models; merges after #2741 [3/3] (#2741). Split from #2774, which first routed every load in one PR and whose non-expert-parallel half is now #2779: this half needs the expert-parallel loader, so it lands after it. Applying the model spec in the decision and the load, which this PR carried before the series was reordered, moved to #2746 with the specs, and came back here when #2746 was folded in on 2026-11-03. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Pretrained MoE models can load with routed experts sharded across ranks under `torchrun`. Distribution plans also support FSDP2 sharding or loading a full copy on each rank. * **Bug Fixes** * Loading now reports an error when a routed-expert model would otherwise remain unsharded. * **Documentation** * Updated loading guidance and changelog details for distribution plans. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
main. The one PR this needs, #2740 (distributed_forward_loop), has merged.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).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: 221 passed; 5.5.4: 206 passed, 15 skipped).hfep-base-6a-ep-loader(after [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 merged): no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 221 passed; 5.5.4: 206 passed, 15 skipped).hfep-base-6a-ep-loader(after [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 merged): no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 221 passed; 5.5.4: 206 passed, 15 skipped).hfep-base-6a-ep-loader(after [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 merged): no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 221 passed; 5.5.4: 206 passed, 15 skipped).main: no conflicts, own diff unchanged. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 230 passed; 5.5.4: 215 passed, 15 skipped).distributed_forward_loop'sforwardis annotated as taking aMapping, since without adeviceit gets the loader's own batch. Its docstring says that only with adevicedoes it get a dict of its own to change. No behaviour change. Pre-commit passes against main; the three test files pass on CPU (transformers 5.18: 230 passed; 5.5.4: 215 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].
Summary by CodeRabbit