Repository navigation
[OMNIML-6060] Load an expert-parallel plan through load_pretrained - #2780
shengliangxu wants to merge 32 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation |
|---|---|---|
| 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 | |
The pull request changes no security-sensitive loading pattern covered by this check. In the only changed modelopt Python file, trust_remote_code remains caller-configurable with a default of `False… |
| Title check | |
The title clearly and concisely describes the main change: routing an expert-parallel plan through load_pretrained. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
The pull request changes no security-sensitive loading pattern covered by this check. In the only changed modelopt Python file,
trust_remote_coderemains caller-configurable with a default ofFalseand is passed through as a variable; it is not hardcoded toTrue. The diff adds notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True),eval,exec, or# nosec. Nopyproject.tomlorrequirements.txtdependency changes are present.
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
|
2f18815 to
d68e7ff
Compare
68f620d to
5c5cde0
Compare
shengliangxu
changed the title
d68e7ff to
da76af0
Compare
5c5cde0 to
c892568
Compare
da76af0 to
646405a
Compare
c892568 to
596bb8f
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-6c-fp8-dequant #2780 +/- ##
================================================================
Coverage 61.86% 61.86%
================================================================
Files 649 649
Lines 73059 73064 +5
================================================================
+ Hits 45195 45200 +5
Misses 27864 27864
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:
|
596bb8f to
8c34d8d
Compare
646405a to
678d634
Compare
shengliangxu
changed the title
|
/ok to test 8c34d8d |
8c34d8d to
d2db2ee
Compare
22676c9 to
df0496c
Compare
06fdac7 to
cd46f22
Compare
|
/ok to test cd46f22 |
load_and_prepare_ep freezes the model it loads unless told otherwise, which the FSDP2 and whole-model loads do not. load_pretrained decides which load runs from the world size and the model, so the same call returned a frozen model on some runs and not on others. Pass freeze=False unless the caller asks for freeze, as the other loads return the model. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…acement With neither a plan nor MODELOPT_EP_DEGREE, the plan decision can shard an MoE's experts, and the expert-parallel load then refuses options that place a whole model, such as the device_map hf_ptq passes, as if the caller had asked for it. Say the layout was decided, and how to keep the load whole, as the FSDP2 load already does when rank 0's GPU memory chose to shard: the measured flag _refuse_sharded_load_options takes now also covers a layout the decision chose for the model's routed experts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The refusal of a whole-model placement on a decided layout read as placement every option the FSDP2 load does not take, freeze and require_lm_head among them. So with neither a plan nor MODELOPT_EP_DEGREE, load_pretrained(..., freeze=True) or require_lm_head=False, the expert- parallel load's own options, was refused as if it placed a whole model. An option either sharded load takes places nothing: _refuse_sharded_load_options, where the placement a measured layout refuses is now computed, leaves both loads' options out of it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The expert-parallel load now refuses a run in which only some ranks have a GPU of their own, so no rank is left waiting in its collectives. The refusal before it still covers a run in which none has one, and names the way to a whole model, which the load's own check does not. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A sharded load refused an option only the other sharded load takes as one that places a whole model, and named the way to a whole model on every rank, where the whole-model load refuses the option again. It now says which load takes it. Which options a sharded load counts as passed now follows the one rule the whole-model load drops them by, _ASKS_NOTHING's: an FSDP2-only option asks nothing at the FSDP2 load's own default, any other at None. So broadcast_chunk_size=None, which asks the FSDP2 load to broadcast all of a source's layers at once and was dropped without a word on the expert-parallel load, is refused as the FSDP2 load's option there, and cpu_offload=False, what hf_ptq passes by default, asks nothing of it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The refusal of an option only the other sharded load takes said which load takes it, but not what to do. It now says to drop the options, or to pass a plan that takes that load: one with ep_size > 1 for the expert-parallel load's, and one with ep_size=1 and dense='fsdp' for the FSDP2 load's, since an expert-parallel plan that FSDP2-shards the dense part is refused too. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A launch that shows each process one GPU, bound to cuda:0, gives every rank a GPU of its own whatever its local rank, and the GPU refusal before the expert-parallel load takes it so through the shared predicate. A case of the expert-parallel load's routing test, on local rank 1 with one GPU visible, holds it: under the predicate's older reading, by local rank alone, that rank is refused. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The backward-breaking entry for a distributed load_pretrained still said an MoE loads without expert parallelism until the expert-parallel load lands. This is that load: an MoE has its routed experts sharded across the ranks, and is refused where it would otherwise load whole on every rank. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With the EP degree left to the plan decision, an MoE whose dense part, replicated beside a rank's share of the routed experts, does not fit one GPU was planned expert-parallel with the dense part FSDP2-sharded, a layout the expert-parallel load refuses, as it replicates the dense part. Such a load, which sharded the whole model with FSDP2 before, now failed. An EP degree the decision chose now gives way to FSDP2 for the whole model there; one the caller asked for stands, and the load refuses it naming the ways out. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With nothing asking for expert parallelism, the plan decision committed any MoE whose config carries a base_model_ep_plan to the expert-parallel load, the ones it refuses included: a plan that shards no expert weights through grouped_gemm, or styles them by one it drops (GPT-OSS's and Llama-4's below their specs), and a checkpoint format it does not load. Its refusals named no way back, and some came only after the download. An EP degree the decision chose now gives way to none there, decided by the load's own checks, which _ep_plan_refusal now holds for both, and such an MoE is not refused for being left whole, as it cannot be expert-parallel here. An EP degree the caller asked for still meets the load's refusal. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The refusal of an expert-parallel plan that FSDP2-shards the dense part said a plan left to decide shards it when a rank's share would not fit, which the decision no longer does: an EP degree it chose gives way to FSDP2 for the whole model there, and only one the caller asked for meets this refusal. Say that, keeping the 60% and the ways out. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
meet the expert-parallel load. A launch with no process group that sets MODELOPT_EP_DEGREE needs only the group, as one that sets MODELOPT_FSDP_SIZE does. The expert-parallel load builds the checkpoint's own class, a VLM's image-text class included, and its refusal of an auto_class names its own way to a whole model. The docstring and the decision's auto_class refusal say which class each sharded load builds. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…y so The plan decision's give-way for a model the expert-parallel load would refuse read the plan's styles and the checkpoint format, but not a model spec ruling expert parallelism out (Llama-4's, GPT-OSS's, whose experts the load refuses once built) nor a transformers below the load's floor or the spec's: an EP degree it chose still committed those to a refusal, after the download. _ep_load_refusal now gathers every reason the config tells, and the decision gives way for any of them. Each give-way now says why on rank 0, as the plan's other outcomes do: the load's reason, or for an FSDP2-sharded dense part the knobs that shard the experts anyway. The changelog names only the MoEs the expert-parallel load takes as refused when left whole. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The decision's give-way and load_and_prepare_ep each spelled 5.7 out, so raising the load's floor in one place would leave the decision planning expert parallelism for a model the load then refuses. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Only an EP degree the decision picked consulted the load's refusal, so an EP degree of 1 the caller set left an MoE the load refuses (GPT-OSS by its spec, a quantized checkpoint, a plan sharding no expert weights) to the left-whole refusal, which offered expert parallelism as the way out. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The load checked a plan's styles alone, so an EP degree asked for, or a plan handed in, loaded a model whose spec rules expert parallelism out (GPT-OSS's, Llama-4's) or states a later transformers, which the decision plans without expert parallelism. The load now asks the same _ep_load_refusal, after its own floor and format checks, which keep their messages. The load tests' stand-in configs carry a model type for the spec lookup, one no spec names. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_ep_load_refusal now returns the error rather than its message: a model spec ruling expert parallelism out and a checkpoint format the load does not take are NotImplementedError, a transformers below the floor a spec states is the RuntimeError require_transformers raised, and a plan's styles stay a ValueError. The load raises it as it is, so a floor refusal has one type whichever floor it is below; the plan decision only prints it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The give-way test tripped only the load's own floor. A case with a spec stating 5.13 on transformers 5.10 decides it by the spec's floor, which a string comparison would put below 5.7. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
#2779's guide, written while no load sharded the routed experts, said a plan with ep_size > 1 is refused. It now lists the expert-parallel route, how the decision picks an EP degree and when it gives way, the way to a whole model on that route, and what the expert-parallel load builds, takes and refuses beside the FSDP2 load. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
#2779 refuses a sharded load inside init_quantized_weights, whose from_pretrained the sharded loads never call. The expert-parallel load meets that refusal too, naming its own way to a whole model: an EP degree of 1 as well, for a plan decided from MODELOPT_EP_DEGREE or handed in. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test c2d18e6 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — LGTM: prior rounds' concerns are resolved and I found no defect in the expert-parallel routing, but two existing tests are removed, so a human signs off.
Baseline: .agents/review-guidelines/code-review.md at base 930bb0f0e1c. Reviewed head c2d18e6c473, full scope (the trigger carried no scoping instructions); all 5 changed files opened, nothing skipped. 351 added source lines, inside the size budget. I ran no tests and report no CI result — findings are from source only.
Needs action:
- Confirm the two removed tests as a maintainer:
test_no_ep_degree_dividing_the_experts_leaves_them_wholeandtest_an_expert_parallel_plan_is_refusedboth covered behaviour this PR replaces, and new rows assert the new outcome — justified from the diff, but a bot cannot auto-approve a test removal. - Add the
MODELOPT_FSDP_SIZE=world+ auto-degree row totest_an_ep_capable_moe_may_be_fsdp2_sharded_instead— only the memory-driven half of the new give-way atmodel_load_utils.py:1654is pinned (inline).
No action needed:
- ✔️ Resolved since the last rounds: the EP route's
init_quantized_weightsrefusal now names its own way to a whole model,_ep_load_refusalreturns the error so a spec's higher transformers floor raises as aRuntimeError,load_and_prepare_epmeets the refusals the decision gives way for, and the loading guide describes the EP route. - Traced clean:
_routeis reached with an identical plan on every rank at every call site;_refuse_a_rank_without_its_gpusits inside the collectiveload checksblock;_decide_distributionruns on rank 0 only andbroadcast_fromships its newValueErrorto every rank — no rank can be stranded in a collective.plan,hub_options,freezeandrequire_lm_headall have live use paths, and the_SHARDED_LOAD_OPTIONS[route]filter keeps strayNoneoptions offload_and_prepare_ep's fixed signature._ckpt_quant_methodreadstext_config/llm_config, so a VLM whose language model is quantized is still refused on the EP route.CHANGELOG.rstamends the same unreleased 0.49.0 entries. - Two smaller suggestions inline: the module-level
get_specimport makes line 802's deferred-import comment stale, and_ep_load_refusal's floor/format branches are dead insideload_and_prepare_ep.
🤖 Generated with Claude Code
| monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") | ||
| monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") |
There was a problem hiding this comment.
[SUGGESTION] The new give-way at model_load_utils.py:1654 fires for two different reasons, and only one of them is pinned.
if chosen and ep_size > 1 and dense == "fsdp" is reached either because the memory heuristic chose "fsdp" (dense_chosen true — covered by test_what_a_rank_holds_with_the_dense_part_replicated_must_fit, whose decided case asserts DistributionPlan(4, dense="fsdp")) or because the caller asked for it with MODELOPT_FSDP_SIZE=world / dense="fsdp" while leaving the EP degree on auto. Nothing covers the second: this test pins the degree to 1 itself, so the auto degree never has to give way, and test_the_env_knobs_override_the_automatic_plan sets the degree explicitly (where it is meant to stand).
That leaves a live mutation gap. Narrowing the condition to chosen and dense_chosen and ... would pass CI, and MODELOPT_FSDP_SIZE=world on an 8-expert MoE over 4 ranks would again decide DistributionPlan(4, 4, "fsdp") — a plan _route refuses outright, so the load would fail where it now succeeds. The else arm of the why message ("is to be sharded") is likewise unexercised.
Cheapest fix is a parametrized row here rather than a new test:
@needs_qwen3_moe_ep_plan
@pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"])
def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(
run_as, serve, monkeypatch, tmp_path, ep_degree
):
"""Asked for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a
replica on every rank: an EP degree left to the decision gives way, as the expert-parallel
load replicates the dense part."""
if ep_degree is not None:
monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree)
monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
run_as(2)
serve(_qwen3_moe(num_experts=8))
assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp")The auto row fails if the dense_chosen reason is the only one honored.
There was a problem hiding this comment.
Taken in f7b11fe. The test now runs with the EP degree left to the decision as well as pinned to 1. Narrowing the give-way to chosen and dense_chosen and ... fails the auto case.
| from transformers.conversion_mapping import get_model_conversion_mapping | ||
| from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key | ||
|
|
||
| from modelopt.torch.models.specs import get_spec |
There was a problem hiding this comment.
[SUGGESTION] This top-level import pulls in the whole modelopt.torch.models package, which makes the deliberately deferred import at line 802 — and its comment — stale.
modelopt/torch/models/__init__.py imports every per-model subpackage for its registration side effect, so from modelopt.torch.models.specs import ... executes the parent __init__ and registers every spec at this module's import time. Line 802 then says:
# Importing modelopt.torch.models registers every model spec; only callers of this need them.
from modelopt.torch.models.hf import prepare_model_for_loadingwhich no longer buys anything — the package is already imported by the time load_pretrained runs. Either hoist prepare_model_for_loading to the top beside get_spec and drop the comment, or import get_spec lazily inside _ep_load_refusal to keep the laziness the comment claims. (I checked for a cycle: nothing under modelopt/torch/models/ imports model_load_utils, so this is import cost and a misleading comment, not a breakage.)
| # The refusals the plan decision gives way for, which an EP degree asked for, or a plan | ||
| # handed in, meets here: a model spec ruling expert parallelism out or stating a later | ||
| # transformers, and the plan's styles. | ||
| if (refusal := _ep_load_refusal(hf_config, ep_plan)) is not None: |
There was a problem hiding this comment.
[SUGGESTION] Two of _ep_load_refusal's four branches are unreachable from here, so the load now checks the same two things twice.
By the time this line runs, load_and_prepare_ep has already called moe.require_transformers(_EP_MIN_TRANSFORMERS, ...) and refused _ckpt_quant_method(hf_config) not in _EP_SOURCE_FORMATS with its fuller "it loads unquantized checkpoints… Start from an unquantized copy" wording. _ep_load_refusal repeats both: its base-floor check can never fire here (only the spec's higher floor can), and its terse f"The expert-parallel load cannot dequantize a {method!r} checkpoint." is dead in this path while being the message the decision prints.
Not a defect — the decision needs all four — but it is two redundant checks on the load path and two messages for the same refusal. Worth either splitting the spec-only part out for the load, or moving this call above the floor/format checks and deleting the duplicates (accepting the terser format message, or moving the fuller text into _ep_load_refusal so both callers say the same thing).
…ed for The decided EP degree gives way to FSDP2 for the whole model either when rank 0's GPU memory chose to shard the dense part or when the caller asked for it. The test pinned the EP degree to 1 itself, so only the first was checked; it now also leaves the degree to the decision. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test f7b11fe |
|
/claude review |
There was a problem hiding this comment.
Claude review — LGTM: the one open action item from the last round is now pinned, and I found no defect in the expert-parallel routing; a human signs off only because existing tests are removed.
Baseline: .agents/review-guidelines/code-review.md at base 930bb0f0e1c. Reviewed head f7b11fea, full scope (the trigger carried no scoping instructions); all 5 changed files opened, nothing skipped. 351 added source lines, inside the size budget. I ran no tests and report no CI result — findings are from source only.
Needs action:
- Confirm the removed/replaced tests as a maintainer:
test_no_ep_degree_dividing_the_experts_leaves_them_whole,test_an_expert_parallel_plan_is_refused,test_a_load_it_decides_keeps_the_experts_wholeand theep-knobmatch intest_what_needs_the_group_is_refused_in_a_launch_that_sets_none_upall pinned behaviour this PR replaces, and new rows assert the new outcome — justified from the diff, but a bot cannot auto-approve a test removal.
No action needed:
- ✔️ Resolved:
test_an_ep_capable_moe_may_be_fsdp2_sharded_insteadnow parametrizesep_degreeover[None, "1"], so the caller-asked half of the give-way atmodel_load_utils.py:1654is pinned alongside the memory-driven half intest_what_a_rank_holds_with_the_dense_part_replicated_must_fit. - Traced clean this round:
_decide_distributionresolvesrefusalbefore the"auto"give-aways and the dense-memory estimate, so the order is self-consistent — and because the left-wholeValueErroris gated onrefusal is None, a quantized MoE too large to replicate now has a working way out (MODELOPT_FSDP_SIZE=1) rather than the dead end an earlier round flagged. The module-scopefrom modelopt.torch.models.specs import get_specimportsmodelopt.torch.models, which registers every spec, soget_specis populated even whenresolve_distribution_planis called directly. - Two earlier suggestions remain open and stay non-blocking: the deferred-import comment at
model_load_utils.py:802is stale, and_ep_load_refusal's format branch is unreachable insideload_and_prepare_ep.
🤖 Generated with Claude Code
Note
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: #2741 [2/3] (#3136). 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.Series layout
mainrequire_transformers, and a VLM built without its LM headhfep-base-6a-ep-loader([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 + [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 + [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 + [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131)hfep-base-9-specs([OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 + [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135)What does this PR do?
Type of change: New feature
#2779 routes
load_pretrainedby aDistributionPlan, 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:load_and_prepare_ep, handed the plan itself. Its arguments are checked with the FSDP2 load's, in_sharded_load_kwargsinside [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's collective "load checks" block;_load_shardedthen runs the load. That load replicates the dense part; FSDP2-sharding it comes with [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742, a sibling of this PR. Itsfreezeandrequire_lm_headoptions pass through, and so do the Hub options (token,cache_dir,local_files_only), which its config read andfrom_pretraineduse. The sharded loads' refusals from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 cover it too: a Nemotron-H checkpoint's MTP head, whichload_and_prepare_epdoes not build, is refused rather than exported unquantized. Until [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #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.MODELOPT_EP_DEGREE=1withMODELOPT_FSDP_SIZE=1, or a replicated plan handed toload_pretrained, which skips the decision. A replication both knobs ask for loads, with a note that every rank holds the whole MoE. So [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's refusals of routed experts the decision can neither count nor shard by now adviseMODELOPT_EP_DEGREE=1withMODELOPT_FSDP_SIZE=1(ep_size=1withdense='replicate'for a caller passing the arguments): an EP degree of 1 alone now meets the refusal above.This closes the series' second phase: an MoE whose expert-parallel plan ships with transformers, such as Qwen3-MoE, now loads expert-parallel through
load_pretrainedand calibrates. The model specs come next (#2746):Usage
Testing
tests/unit/torch/utils/test_model_load_utils.py(where [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #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_wholegoes: the refusal replaces it. [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's two refusal tests now setMODELOPT_FSDP_SIZE=1in theirset-to-1rows and check the refusals name it.tests/unit/torch/utils/test_load_pretrained_distribution.py([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779) gains the load's cases:load_and_prepare_epwith the plan, records the plan, and goes through no other loader;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.pygains a Qwen3-MoE case: left to decide,load_pretrainedshards 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/uniton 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 [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #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 [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 (was on [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743), with three follow-ups:
test_resolve_distribution_plan.py, which [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 deleted, intotest_model_load_utils.py;MODELOPT_FSDP_SIZE=1too;load_pretraineddocstring 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.pyandtest_distributed.py.Review fix, 2026-10-29, from the re-review after the re-cut: on [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 the expert-parallel load refuses a sharded dense part, so
_routenow 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 [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [2/3] ([OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136); the former head is pinned.
_sharded_load_kwargs, inside the collective "load checks" block, and its load in_load_sharded.test_load_pretrained_distribution.pyfiles moved intotest_model_load_utils.py(unit and GPU), where [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 moved the load tests._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):device_map), it names its way to a whole model,MODELOPT_EP_DEGREE=1withMODELOPT_FSDP_SIZE=1or 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.pyandtest_ep_loader.pypass on transformers 5.18 (161) and 5.5.4 (139, 22 skipped).test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (405) and 5.5.4 (376, 29 skipped).its-own-plan) oftest_an_expert_parallel_moe_loads_with_its_experts_shardedat [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 [3/3]'s ae518d6, passed. That version of the test also compares each rank's logits with the whole model's. [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746's spec-plan row and the two denseload_pretrainedtests passed beside it, 4 tests in all. NVIDIA GB200, torch 2.11.0a0, transformers 5.9.0, reruns disabled:_load_shardedtakes the route beside [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's new_has_own_gpu.load_pretrained's docstring: the expert-parallel bullet goes before [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's reworded FSDP2 bullet, and the dropped "anep_size > 1plan raises" sentence leaves the rest of [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779's paragraph as it is.CHANGELOG.rstline keeps this PR's wording.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (467) and 5.5.4 (425, 31 skipped).mp_policyfor the FSDP2 load only (90c914a).test_a_sharded_load_refuses_what_it_cannot_honorgains ano_gpurow and anmp_policyassertion. Each fails without its change.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (469) and 5.5.4 (427, 31 skipped).freeze,require_lm_head) before anything loads, naming that load. Before, they reachedfrom_pretrainedand failed inside transformers.test_a_load_refuses_what_it_cannot_loadgainsep-option-on-a-whole-load, which fails without the change.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (470) and 5.5.4 (428, 31 skipped).ep_size/densearguments, which override the environment knobs, or the knobs (b1f5d8c). A newexplicit_degree_one_as_an_argumentrow fails without it.freeze=False, unlike the other loads (1072e8f).test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (471) and 5.5.4 (428, 32 skipped).ValueError, like the decision's other refusals of what its inputs resolve to (d0a63ff).test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (471) and 5.5.4 (428, 32 skipped).freezeorrequire_lm_headpassed asNoneasks nothing, as other options passed asNonedo, rather than reading as False and dropping the LM head (83d9ce0).freeze=True, as the other routes do (69e56a9).test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (472) and 5.5.4 (428, 33 skipped).freezeandrequire_lm_head, rather than refusing them as a whole-model placement (1db4cc3). A newits-own-optionsrow fails without the change._refuse_a_rank_without_its_gpusays what the check adds beside [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's GPU check in the expert-parallel load, which now refuses a run where only some ranks have a GPU (8f9c8d9).test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (476) and 5.5.4 (430, 35 skipped).broadcast_chunk_size=Noneasks 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 oftest_a_load_refuses_what_it_cannot_loadeach fail without the change.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (478) and 5.5.4 (432, 35 skipped).test_a_load_it_decides_shards_the_expertsreplaces [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779'stest_a_load_it_decides_keeps_the_experts_whole._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.measuredrefusal, which this PR extends to that layout. Only options that neither sharded load takes count as placement._ASKS_NOTHING. On the expert-parallel load,broadcast_chunk_size=Noneis refused as an FSDP2 option, andcpu_offload=Falseasks for nothing.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (513) and 5.5.4 (463, 39 skipped).test_what_a_rank_holds_with_the_dense_part_replicated_must_fitasserts both outcomes, and its decided case fails without the change.load_and_prepare_epdocstring and comments, plus this PR's sentence onhub_options.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (515) and 5.5.4 (463, 41 skipped)._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_itfails without the change in both of its cases.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (519) and 5.5.4 (464, 44 skipped).its-specandits-transformerscases fail without the change.test_model_load_utils.py,test_ep_loader.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (521) and 5.5.4 (464, 46 skipped).test_ep_loader.pyintotest_model_load_utils.py. This PR's one change to that file, the Hub-options checks intest_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.pyandtests/unit/torch/modelspass on transformers 5.18 (521) and 5.5.4 (464, 46 skipped)._EP_MIN_TRANSFORMERS, that the decision and the load both read (6210606).MODELOPT_EP_DEGREE=1, and theset-to-1cases fail without the change.test_model_load_utils.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (525) and 5.5.4 (464, 50 skipped).load_and_prepare_eprefuses 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.test_model_load_utils.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (526) and 5.5.4 (465, 50 skipped).init_quantized_weightsrefusal 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 [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #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 inload_pretrained._ep_load_refusalreturns the error itself, so a floor refusal is aRuntimeErrorwhichever floor it trips (eda4c88). A spec's floor above the load's own is tested (608b0b8). Both cases fail without their change.test_model_load_utils.py,test_moe_a2a.pyandtests/unit/torch/modelspass on transformers 5.18 (536) and 5.5.4 (471, 54 skipped).Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional 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.
Summary by CodeRabbit
torchrun. Distribution plans also support FSDP2 sharding or loading a full copy on each rank.