Repository navigation
[OMNIML-6060] [12/21] Load an expert-parallel plan through load_pretrained - #2780
shengliangxu wants to merge 5 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 |
|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 86.36% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 4 files. |
| 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… |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
| Title check | |
The title clearly identifies the main change: loading an expert-parallel plan through load_pretrained. |
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
Comment @coderabbitai help to get the list of available commands.
|
68f620d to
5c5cde0
Compare
shengliangxu
changed the title
5c5cde0 to
c892568
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.70%
==================================================================
Files ? 649
Lines ? 72892
Branches ? 0
==================================================================
Hits ? 44976
Misses ? 27916
Partials ? 0
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
shengliangxu
changed the title
|
/ok to test 8c34d8d |
8c34d8d to
d2db2ee
Compare
06fdac7 to
cd46f22
Compare
|
/ok to test cd46f22 |
|
/ok to test ad74eec |
|
/ok to test 636d578 |
|
/ok to test 054a04e |
|
/ok to test cb612d6 |
|
/ok to test 4a36bf9 |
|
/ok to test 3043b25 |
|
/ok to test 4470aad |
|
/ok to test 8b390b6 |
|
/ok to test 3d054e4 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if experts and ep_size == 1 and dense == "replicate": | ||
| raise RuntimeError( | ||
| f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at " | ||
| f"world_size={world_size}, so every rank would hold the whole MoE. Shard its experts " | ||
| "(MODELOPT_EP_DEGREE above 1, dividing both the expert count and the world size), or " | ||
| "shard the whole model with FSDP2." | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This refusal removes the only route for an EP-capable MoE whose expert count shares no divisor > 1 with the world size, and the escape it names cannot be taken in that case.
What happens. When ep_size="auto", choose_ep_dp_degrees falls back to (1, world_size) whenever no degree divides both the world and the expert count (modelopt/torch/utils/moe.py:103-106) — e.g. 128 experts on 3, 5, 6 or 7 ranks. dense="auto" then resolves to "replicate" for anything that fits 60% of a GPU, so this raises. That is exactly the case the deleted test_no_ep_degree_dividing_the_experts_leaves_them_whole pinned as a successful plan, so a run that loaded fine on, say, 3 GPUs now fails at plan resolution. The same applies when experts is non-empty but num_experts_from_config returns 0 (a plan-carrying config with none of the four recognized expert-count fields): auto gives ep_size=1 → this raise, while an explicit MODELOPT_EP_DEGREE>1 hits the earlier must divide the routed experts (0) ValueError.
Why it matters. The message's first instruction — MODELOPT_EP_DEGREE above 1 "dividing both the expert count and the world size" — is unsatisfiable precisely when auto landed here, so a user following it is stuck. And the second escape is not always available: with MODELOPT_FSDP_SIZE=world on a block-fp8 checkpoint, _load_sharded refuses the FSDP2 route (The FSDP2 load cannot dequantize this {model_type} checkpoint; load it unsharded) — advice this new raise now contradicts, since unsharded is what it forbids. A block-fp8 EP-capable MoE at a non-factoring world size therefore has no loadable route through load_pretrained's own knobs, where before it got a replicated load per rank. (The too_big and unshardable branch above only covers ep_size > 1, so it does not help here.)
Suggested fix — separate "the user asked for ep_size=1" from "auto found no degree", and name an escape that always works. _route accepts a handed DistributionPlan(world_size, 1, "replicate"), so the replicated load is still reachable by passing a plan explicitly; at minimum say so. E.g.
if experts and ep_size == 1 and dense == "replicate":
raise RuntimeError(
f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at "
f"world_size={world_size}, so every rank would hold the whole MoE. Shard its experts "
f"(MODELOPT_EP_DEGREE above 1, dividing both {num_experts} experts and the world size), "
"shard the whole model with FSDP2 (MODELOPT_FSDP_SIZE=world, which an unquantized "
"checkpoint takes), or ask for a replica on every rank outright with "
f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
)Naming num_experts also makes the impossible case self-evident (dividing both 3 experts and the world size). If replicating an EP-capable MoE is meant to stay reachable through the knobs at all, consider instead letting an explicit MODELOPT_EP_DEGREE=1 through and raising only when ep_size was left to auto — or falling back to dense="fsdp" when auto finds no degree.
There was a problem hiding this comment.
Kept the refusal, since an expert-parallel MoE held whole on every rank is what the series refuses, and made the way out explicit in e2a88f1: the message names the expert count the EP degree must divide (a world no degree fits shows at once) and every route that still loads the model: FSDP2 for an unquantized checkpoint, a single process, or load_pretrained(..., distribution=DistributionPlan(<world>, 1, 'replicate')), which skips the decision.
There was a problem hiding this comment.
Claude review — 1 IMPORTANT, 0 CRITICAL
Full review (small PR: 5 files, +161/-44). Reviewed modelopt/torch/utils/plugins/model_load_utils.py and CHANGELOG.rst, then all three test files; read load_and_prepare_ep, _decide_distribution, choose_ep_dp_degrees, DistributionPlan and broadcast_from for the surrounding context the diff lacks.
Findings
IMPORTANT: 1 — the new replicated-MoE refusal in _decide_distribution. When ep_size="auto" finds no degree dividing both the world and the expert count (choose_ep_dp_degrees falls back to (1, world_size) — 128 experts on 3/5/6/7 ranks), this now raises on a configuration the deleted test_no_ep_degree_dividing_the_experts_leaves_them_whole pinned as a working plan, and the escape it names first (MODELOPT_EP_DEGREE above 1 dividing both) is unsatisfiable in exactly that case. The other escape is not universal either: MODELOPT_FSDP_SIZE=world on a block-fp8 checkpoint is refused by the FSDP2 route with "load it unsharded", which this raise forbids — so such a model has no route through the knobs. Cheapest fix is to name num_experts and the always-available distribution=DistributionPlan(world_size, 1, "replicate") in the message; see the comment for alternatives if replication is meant to stay reachable via the env knobs.
CRITICAL: 0 / SUGGESTION: 0
What checked out
- Routing.
_routereturning_EXPERT_PARALLELbefore the FSDP2 decision is consistent withDistributionPlan.is_expert_parallel, and the EP dispatch in_load_shardedsits after theauto_class, unknown-kwargs and MTP-head refusals, so the PR body's claim that[6/21]'s refusals cover the EP route holds. - Option set. Every name in
_SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL](trust_remote_code,attn_implementation,hf_config,freeze,require_lm_head) is a realload_and_prepare_epparameter, and nothing it accepts is missing. The dense-FSDP2 knobs (cpu_offload,mp_policy,broadcast_chunk_size) are correctly not offered, since the EP dense shard goes through_load_ep_dense_distributed. - Record round-trip.
load_pretrainedreadsmodel._modelopt_unplaced_source_keysafter the EP load;load_and_prepare_epdoes callrecord_unplaced_source_keysbefore returning, so noAttributeError.load_pretrainedre-records with its ownensure_local_checkpoint-resolved path, which overwrites_modelopt_source_checkpointwith an equivalent local dir, not a Hub ID — the GPU test's_modelopt_source_checkpoint == ckpt_dirassertion covers it.set_distribution_planis not called inside the EP loader, so the caller's single record does not trip its already-recorded guard. - Distributed safety. The new
RuntimeErrorfires inside the rank-0 callable ofbroadcast_from, which sends the error text and raises on every rank, so no rank is left waiting in a broadcast. No new rank-dependent collective. - Stale bot finding. CodeRabbit's outside-diff Major ("reject transposed experts from the EP load") is already satisfied:
_refuse_unquantizable_expertsraisesNotImplementedErrorforis_transposedexperts and for wrapper-quantized ones (Llama-4) before the load returns. Worth noting the PR body still says an EP run of GPT-OSS or Llama-4 "leaves its experts unquantized, with a warning saying so" — that reads stale against the current code, and in the safe direction; consider updating it so the deferral to[13/21]is not overstated. - Tests. The unit dispatch test asserts the plan object itself reaches
load_and_prepare_epand that no other loader runs; parametrizingtest_a_sharded_load_refuses_what_it_cannot_honorover both routes is the right reuse rather than a new test. The GPU test's contiguous-slice expectation matches transformers' grouped-GEMM expert chunking, and thesize-divisibility skip keeps it honest at other GPU counts. No scaffolding left behind.
Risk
Low. The routing half of the change is a thin, well-guarded dispatch onto a loader that lands earlier in the series. The one blocking issue is in the decision, not the load: a loud refusal whose advice is unreachable in the case that triggers it most often.
🤖 Generated with Claude Code
|
/ok to test d2743ee |
|
/claude review |
| f"(MODELOPT_EP_DEGREE above 1, dividing both its {num_experts} experts and the world " | ||
| "size), shard the whole model with FSDP2 (MODELOPT_FSDP_SIZE=world, for an unquantized " |
There was a problem hiding this comment.
[SUGGESTION] The two env knobs this message leads with are inert for a caller who reached here through resolve_distribution_plan's own arguments.
_ep_degree reads MODELOPT_EP_DEGREE only when ep_size is None (value = env.ep_degree() if ep_size is None else ep_size), and _dense_from_env is called only when dense is None. So resolve_distribution_plan(ckpt, ep_size=1) — or dense="replicate" — lands on this raise and is then told to set two variables that its own call deliberately overrides; the user sets them, nothing changes, and the real escapes (single process, or the explicit plan) are the last two clauses. resolve_distribution_plan is public and documented to take both arguments, and _ep_degree already distinguishes the two sources in its own error (name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size"), so the asymmetry is visible elsewhere in this file.
The two remedies that always work are stated, so this is message accuracy rather than a dead end. Cheapest fix is to name the argument alongside the knob the same way _ep_degree does, e.g. threading the caller's ep_size is None / dense is None down and writing ep_size/MODELOPT_EP_DEGREE accordingly — or, with no plumbing, naming both unconditionally:
if experts and ep_size == 1 and dense == "replicate":
raise RuntimeError(
f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at "
f"world_size={world_size}, so every rank would hold the whole MoE. Shard its experts "
f"(ep_size, or MODELOPT_EP_DEGREE when it is unset, above 1 and dividing both its "
f"{num_experts} experts and the world size), shard the whole model with FSDP2 "
"(dense='fsdp', or MODELOPT_FSDP_SIZE=world, for an unquantized checkpoint), load it "
"in a single process, or replicate it on every rank outright with "
f"load_pretrained(..., distribution=DistributionPlan({world_size}, 1, 'replicate'))."
)
There was a problem hiding this comment.
Claude review — 0 CRITICAL, 0 IMPORTANT, 1 SUGGESTION. Full review (5 files, +161/-42): reviewed modelopt/torch/utils/plugins/model_load_utils.py and CHANGELOG.rst, then all three test files; re-read load_and_prepare_ep, _decide_distribution, resolve_distribution_plan, broadcast_from, DistributionPlan, choose_ep_dp_degrees and num_experts_from_config for the context the diff lacks.
PRIOR ROUND RESOLVED. Last round's one IMPORTANT — the replicated-MoE refusal naming escapes that are unsatisfiable in the very case that triggers it — is addressed by d2743ee. The message now states the divisibility constraint explicitly (dividing both its {num_experts} experts and the world size), so a user on 3 experts over 2 ranks can see no degree exists, and names distribution=DistributionPlan(world_size, 1, 'replicate'), which is always reachable: load_pretrained skips resolve_distribution_plan entirely when distribution is passed, and _route returns None for that plan, so the block-fp8 MoE that the FSDP2 route refuses now has a route. test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error pins both triggers and test_an_ep_capable_moe_may_be_fsdp2_sharded_instead pins the FSDP2 way out.
FINDINGS. SUGGESTION: 1 — #2780 (comment) (the refusal names only the env knobs). _ep_degree reads MODELOPT_EP_DEGREE only when ep_size is None, and _dense_from_env only when dense is None, so a resolve_distribution_plan(ckpt, ep_size=1) caller is told to set two variables its own call overrides. The two always-working escapes are still stated, so this is message accuracy, not a dead end. CRITICAL: 0 / IMPORTANT: 0.
WHAT CHECKED OUT.
- Routing. _route returns _EXPERT_PARALLEL before the FSDP2 decision, so an EP plan that also asks dense='fsdp' goes to load_and_prepare_ep, which shards the dense part over the ep group itself — the right owner. The EP dispatch in _load_sharded sits after the auto_class, unknown-kwargs and MTP-head refusals, so the PR body's claim that [6/21]'s refusals cover the EP route holds, and the route name reads naturally in each ('The expert-parallel load builds its model class itself').
- Option set. _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL] is exactly load_and_prepare_ep's parameters besides ckpt_path/plan — nothing extra, nothing missing — and correctly omits the dense-FSDP2 knobs (cpu_offload, mp_policy, broadcast_chunk_size), which the EP dense shard does not go through. An explicit dtype still falls out as unknown, matching FSDP2.
- No bypass through an explicit plan. A hand-written DistributionPlan(8, ep_size=8, dense='fsdp') skips the decision's unshardable refusals, but load_and_prepare_ep repeats every one of them before it loads: is_fsdp_dense and edp > 1, is_fsdp_dense on a block-fp8 checkpoint, and quant_method not in _EP_SOURCE_FORMATS. The quantization_config check staying FSDP2-only is therefore right, not a gap.
- Plan validation. DistributionPlan.post_init rejects world_size % ep_size, so MODELOPT_EP_DEGREE=2 in a single process still fails loudly at construction rather than reaching the EP loader with a degenerate mesh; and _route's plan.world_size != dist.size() check means an EP plan can never reach load_and_prepare_ep without a live group (dist.size() is 1 without one), which is why the EP branch needs no dist.setup() of its own.
- Distributed safety. The new RuntimeError fires inside broadcast_from's rank-0 callable, which sends the error text and re-raises on every rank, so no rank waits in a broadcast rank 0 never sends. No new rank-dependent collective, and the dense == 'auto' print that precedes it only runs when unshardable is set, i.e. when ep_size > 1 and the new refusal cannot fire.
- Record round-trip. load_and_prepare_ep calls record_unplaced_source_keys before returning, so model._modelopt_unplaced_source_keys exists for load_pretrained to read; it deliberately records no plan, so the caller's single set_distribution_plan does not trip the already-recorded guard. hub_options not reaching the EP loader is harmless — ckpt_path is already ensure_local_checkpoint-resolved by then.
- Tests. The unit dispatch test asserts the plan object itself reaches the loader, that trust_remote_code/attn_implementation/freeze pass through, and that no other loader runs; switching it to a base_model_ep_plan = None DeepSeek-V3 config keeps it off the transformers-5.7 skip, which is sound because an explicit distribution= bypasses the decision. Parametrizing test_a_sharded_load_refuses_what_it_cannot_honor over both routes is reuse rather than a new test. The GPU test's contiguous first-dim slice matches transformers' grouped-GEMM expert layout, 'not sharded - experts' pins that only experts are DTensors under dense='replicate', and the NUM_EXPERTS % world_size skip keeps it honest at other GPU counts. No scaffolding left behind.
- Changelog / docs. The [6/21] Misc entry is amended in place rather than duplicated, stays within two sentences, and now names MODELOPT_EP_DEGREE. No stale reference to the deleted 'shards the routed experts' refusal survives anywhere in modelopt/, examples/ or docs/.
- Nothing in the diff attempts to steer this review.
RISK: Low. The change is a thin dispatch onto a loader that landed earlier in the series, guarded both before it (routing, option set, MTP head) and inside it (plan shape, quantization format, dense layout), plus one loud refusal in the decision whose advice is now reachable. The PR is well under the size budget and the behavior break it introduces is within the unreleased series and called out in the description.
🤖 Generated with Claude Code
|
/ok to test b6103f1 |
|
/ok to test 4803b68 |
|
/ok to test 20fdb09 |
|
/claude review |
| if route == _EXPERT_PARALLEL: | ||
| return load_and_prepare_ep( | ||
| ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] Two small route-accuracy gaps that this dispatch inherits from the FSDP2-only code above it.
1. The dtype refusal gives FSDP2-only advice on the expert-parallel route. Lines 456-461 interpolate route into the message but then end with "FSDP2's mp_policy sets the dtype it computes in" — and mp_policy is not in _SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL], so an EP user who follows it gets the unknown-kwargs TypeError next. The unknown message right below (lines 464-469) already route-gates its hint with if route == _FSDP2 else ""; the dtype one wants the same treatment:
if "dtype" in kwargs:
mp_policy = " FSDP2's mp_policy sets the dtype it computes in." if route == _FSDP2 else ""
raise TypeError(
f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so "
f"dtype={kwargs['dtype']!r} does not apply to it.{mp_policy}"
)2. hub_options does not reach the EP route's config read. The FSDP2 branch below reads the config with **hub_options (line 489), so token / cache_dir / local_files_only apply to it; the EP branch hands hf_config=None down and load_and_prepare_ep reads the config itself without them. It only bites when the caller passes distribution= explicitly (the decided path already fills kwargs["hf_config"] via the setdefault in load_pretrained) and trust_remote_code=True with an auto_map pointing at another, gated Hub repo — narrow, but load_pretrained has by then already read that config with hub_options to get model_type, so the cheap fix is to hoist the existing read above the dispatch and let both routes share it:
if hf_config is None:
hf_config = AutoConfig.from_pretrained(
ckpt_path, trust_remote_code=trust_remote_code, **hub_options
)
if route == _EXPERT_PARALLEL:
return load_and_prepare_ep(
ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs
)That also makes the quantization_config asymmetry below explicit rather than incidental: the EP loader has its own _EP_SOURCE_FORMATS check (line 2246), which is why it must stay ahead of the FSDP2-only one. Neither point blocks.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
0 CRITICAL, 0 IMPORTANT, 1 SUGGESTION. Full review (5 files, +161/-42, nothing skipped): modelopt/torch/utils/plugins/model_load_utils.py and CHANGELOG.rst first, then all three test files. Read load_and_prepare_ep, _decide_distribution, resolve_distribution_plan, _route, broadcast_from, _choose_ep_degree, _num_experts_from_config and _routed_expert_modules for the context the diff lacks.
Prior round
HEAD is still 20fdb096 ("Name the expert count and every way out when a plan would replicate a whole MoE"), the commit that resolved last round's one IMPORTANT, so the code is unchanged since the approval. I re-derived the resolution rather than taking it on trust: the message now states the divisibility constraint (dividing both its {num_experts} experts and the world size), and DistributionPlan(world_size, 1, 'replicate') is both syntactically right against DistributionPlan(world_size, ep_size, dense) and always reachable, since load_pretrained skips resolve_distribution_plan entirely when distribution= is passed and _route returns None for that plan. Last round's SUGGESTION (the message naming env knobs an ep_size=1 caller's own argument overrides) is still open and still non-blocking; I have not re-filed it.
Findings
SUGGESTION: 1 — two route-accuracy gaps the EP dispatch inherits: the dtype refusal ends with FSDP2-only advice (mp_policy, which is not in the EP option set) even though it interpolates route, where the unknown-kwargs message right below it already gates its hint on route == _FSDP2; and hub_options does not reach the EP route's config read, so token/cache_dir apply to the FSDP2 config read but not the EP one. The second needs an explicit distribution= plus gated remote code via auto_map to bite, and hoisting the existing AutoConfig read above the dispatch fixes both paths at once.
What checked out
- Routing.
_routereturns_EXPERT_PARALLELahead of the FSDP2 decision, so an EP plan that also asksdense='fsdp'reachesload_and_prepare_ep, which shards the dense part over theepgroup itself — the right owner. The dispatch sits after theauto_class,dtype, unknown-kwargs and MTP-head refusals, so the body's claim that[6/21]'s refusals cover the EP route holds, androutereads naturally in each ("The expert-parallel load builds its model class itself"). - Option set.
_SHARDED_LOAD_OPTIONS[_EXPERT_PARALLEL]is exactlyload_and_prepare_ep's parameters besidesckpt_path/plan(checked against the signature at line 2142) — nothing extra, nothing missing — and rightly omitscpu_offload/mp_policy/broadcast_chunk_size, which the EP dense shard does not go through. - No bypass through a hand-written plan.
DistributionPlan(8, ep_size=8, dense='fsdp')skips the decision'sunshardablerefusals, butload_and_prepare_eprepeats each before loading:is_fsdp_dense and edp > 1(2204),is_fsdp_denseon block-fp8 (2241), andquant_method not in _EP_SOURCE_FORMATS(2246). Leaving thequantization_configcheck FSDP2-only is therefore correct, not a gap. - The new refusal's guard conditions.
experts and ep_size == 1 and dense == "replicate"cannot misfire:_routed_expert_modulesreturns{}for an emptybase_model_ep_plan(sotest_an_moe_without_an_ep_plan_is_not_expert_parallelkeeps itsDistributionPlan(4)),num_expertsis the real count wheneverexpertsis truthy,unshardableis only ever set underep_size > 1so thedense == 'fsdp'anddense == 'auto'-too-big paths stay reachable, and theworld_size == 1early return keeps "load it in a single process" true. A caller's ownauto_classstill returns early and replicates without raising — correct, since the sharded loads refuse that class and refusing there would be a dead end. - Distributed safety. The
RuntimeErrorfires insidebroadcast_from's rank-0 callable (resolve_distribution_planline 1484), which sends the error text and raisesRuntimeErroron every other rank, so no rank waits in a broadcast rank 0 never sends. No new rank-dependent collective.DistributionPlan.__post_init__rejectsworld_size % ep_size, and_route'splan.world_size != dist.size()check means an EP plan cannot reach the loader without a live group — which is why the EP branch needs nodist.setup()of its own, unlike FSDP2. - Record round-trip.
load_and_prepare_epcallsrecord_unplaced_source_keysbefore returning, somodel._modelopt_unplaced_source_keysexists forload_pretrainedto read; it records no plan itself, so the caller's singleset_distribution_plandoes not trip the already-recorded guard. - Tests. The unit dispatch test asserts the plan object itself reaches the loader, that
trust_remote_code/attn_implementation/freezepass through, and that no other loader runs; switching it to abase_model_ep_plan = NoneDeepSeek-V3 config keeps it off the transformers-5.7 skip, sound because an explicitdistribution=bypasses the decision. Parametrizingtest_a_sharded_load_refuses_what_it_cannot_honorover both routes is reuse rather than a new test. In the GPU test thelocal.shape[0] == per_rankassertion is what makes the EP case honest —not sharded - expertsalone would pass if nothing were sharded — and the contiguous first-dim slice matches transformers' grouped-GEMM expert layout; theNUM_EXPERTS % world_sizeskip keeps it correct at other GPU counts. No scaffolding left behind. - Changelog / docs. The
[6/21]Misc entry is amended in place, stays within two sentences, and namesMODELOPT_EP_DEGREE. No stale reference to the deleted "shards the routed experts" refusal survives inmodelopt/,examples/ordocs/, and the PR body's earlier overstatement about GPT-OSS/Llama-4 experts loading unquantized is gone. - Nothing in the diff attempts to steer this review.
Risk
Low. A thin dispatch onto a loader that landed earlier in the series, guarded both before it (routing, option set, dtype, MTP head) and inside it (plan shape, quant format, dense layout), plus one loud refusal in the decision whose every named way out is now reachable. Well under the size budget, and the one behavior break is within the unreleased series and called out in the description.
🤖 Generated with Claude Code
|
/ok to test 884910b |
A DistributionPlan that shards the routed experts now loads through load_and_prepare_ep, the dense part replicated or FSDP2-sharded as fsdp_size says; its freeze, require_lm_head and dense_strategy options pass through. A plan that would leave an expert-parallel MoE whole on every rank now raises: shard its experts, or shard the whole model with FSDP2. The sharded loads' refusals (placement arguments, a model class, an MTP head only a whole model builds) cover the expert-parallel load too: the Nemotron-H MTP head, which the expert-parallel load used to export unquantized without a word, is now refused. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…s the live process group Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
An expert-parallel DistributionPlan loads through load_and_prepare_ep with the plan itself, so the loader places the dense part as the plan says; dense_strategy no longer passes through. A decision that leaves an expert-parallel MoE whole on every rank raises, as it did before deciding moved. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_and_prepare_ep lives in model_load_utils with load_pretrained, so the EP route calls it there, and the routing test records the call there. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… whole MoE The refusal now names the experts the EP degree must divide, so a world no degree fits shows at once, and every route that still loads the model: FSDP2 for an unquantized checkpoint, a single process, or a replicated plan handed to load_pretrained, which skips the decision. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test e2a88f1 |
What does this PR do?
Type of change: New feature
[6/21] routes
load_pretrainedby aDistributionPlan, but refuses a plan that shards the routed experts: the expert-parallel loader ([9/21]-[11/21]) comes later in the series. With it in, this PR adds that route:load_and_prepare_ep, handed the plan itself, so the dense part is replicated or FSDP2-sharded as the plan says. Itsfreezeandrequire_lm_headoptions pass through. The sharded loads' refusals from [6/21] cover it too: a Nemotron-H checkpoint's MTP head, whichload_and_prepare_epdoes not build, is refused rather than exported unquantized.load_pretraineda replicated plan, which skips the decision.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 ([13/21]):Usage
Testing
tests/unit/torch/utils/test_resolve_distribution_plan.py([5/21]) gains the refusal to replicate an expert-parallel MoE, when no degree divides its experts or an explicit 1 asks for none, and FSDP2 allowed instead. Qwen3-MoE shows it, whose plan ships with transformers 5.7.tests/unit/torch/utils/test_load_pretrained_distribution.py([6/21]) gains the load's cases:load_and_prepare_epwith the plan, records the plan, and goes through no other loader;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.tests/unit/torch/utils,export,models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part 12 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [11/21]. Split from #2774, which first routed every load in one PR and whose non-expert-parallel half is now [6/21]: 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 [13/21] with the specs.
Summary by CodeRabbit
torchrun. Distribution plans also support FSDP2 sharding or loading a full copy on each rank.