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.62%
==================================================================
Files ? 649
Lines ? 72814
Branches ? 0
==================================================================
Hits ? 44868
Misses ? 27946
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 a1a6fa8 |
|
/ok to test f8b78ed |
|
/ok to test aafc99c |
|
/ok to test 174a19b |
|
Review-body items:
|
|
/ok to test be0b431 |
|
/ok to test a401f63 |
|
/ok to test 7800420 |
|
/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 d2743ee: 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
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 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
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.