Repository navigation
[OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec - #3135
shengliangxu wants to merge 10 commits into
Conversation
|
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 |
|---|---|---|
| Description Check |  |
Check skipped - CodeRabbit’s high-level summary is enabled. |
| Title check | |
The title names the main change: preparing a model's Hugging Face config for expert parallelism from its EP spec. The PR adds apply_distributed_support and helpers that attach the spec's plan to the c… |
| Docstring Coverage | |
Docstring coverage is 90.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati… |
| 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 | |
No listed security anti-pattern was introduced. The only modelopt package change is modelopt/torch/utils/moe.py; no examples, pyproject.toml, or requirements files changed. Added code contains no … |
Full details: Title check
-
Explanation
The title names the main change: preparing a model's Hugging Face config for expert parallelism from its EP spec. The PR adds apply_distributed_support and helpers that attach the spec's plan to the config the load reads, then run the spec's prepare hook. The title matches this primary change. The "[2/3]" tag refers to the stacked PR sequence and does not add noise that hides the change. The title does not mention the support checks or the Transformers version floor, but it does not need to cover every detail.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 90.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 10 functions across 2 files.
Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The only modelopt package change is
modelopt/torch/utils/moe.py; no examples,pyproject.toml, or requirements files changed. Added code contains notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True),eval,exec, or# nosec.trust_remote_codeis an exposed parameter defaulting toFalse, not hardcoded toTrue.
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
|
/ok to test c38ccce |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " | ||
| f"{without_ep}." | ||
| ) | ||
| from packaging.version import parse |
There was a problem hiding this comment.
[SUGGESTION] parse is already imported at module top (modelopt/torch/utils/moe.py:32, used by require_transformers), so this local import is redundant — it rebinds the same name with no effect. CONTRIBUTING's Keep imports at the top of the file allows a function-level import only for a named reason (circular import, optional dependency, heavy import), none of which applies to packaging, which is already a hard import here.
| from packaging.version import parse |
For contrast, the two function-level imports that are justified — from modelopt.torch.models import get_spec in _ep_spec and from transformers... in _routed_expert_modules — the latter carries its reason as a comment; _ep_spec's does not. Worth a one-liner there too (importing modelopt.torch.models pulls in every per-model subpackage for its registration side effect, which this module shouldn't do at import time).
There was a problem hiding this comment.
Taken in 5648c93. check_distributed_support no longer imports parse locally; it uses the module-level import. _ep_spec's local import now says why it is local: importing modelopt.torch.models registers every model spec, which importing this module should not do.
| @@ -58,3 +83,211 @@ def test_the_omni_vlm_puts_its_plan_on_the_nested_config(): | |||
| """Its experts live in the language sub-model, and the globs are written against that tree.""" | |||
| assert get_spec("nemotron_h_omni").ep_spec.plan_config_attr == ("text_config", "llm_config") | |||
| assert get_spec("nemotron_h").ep_spec.plan_config_attr is None | |||
There was a problem hiding this comment.
[SUGGESTION] The plan_config_attr ⊆ _SUB_CONFIG_ATTRS invariant is enforced only at runtime inside apply_distributed_support, and register() does no validation of it (modelopt/torch/models/specs.py:320-330 only checks for duplicate model types). So a new spec that writes plan_config_attr="language_config" or misspells "text_configs" passes every test in this file and only blows up when someone actually loads that model multi-rank with EP — which CI never does.
The applier's own test (test_a_plan_attached_where_the_load_does_not_read_one_is_an_error) covers the mechanism via a synthetic _probe spec, but not the registered data. A one-line data assertion alongside the other EP_MODEL_TYPES spec tests would close that gap cheaply:
@pytest.mark.parametrize("model_type", EP_MODEL_TYPES)
def test_a_plan_goes_where_the_load_reads_one(model_type):
"""A spelling the load does not read is read as no plan at all, and the MoE planned dense."""
attrs = get_spec(model_type).ep_spec.plan_config_attr or ()
attrs = (attrs,) if isinstance(attrs, str) else attrs
assert not set(attrs) - set(moe._SUB_CONFIG_ATTRS), f"{model_type}: {attrs}"This is also the kind of check that stays correct if _SUB_CONFIG_ATTRS ever grows a third spelling, since both sides read the same tuple.
There was a problem hiding this comment.
Taken in 5648c93. test_a_plan_goes_where_the_load_reads_one checks every registered spec's plan_config_attr against _SUB_CONFIG_ATTRS, so a spelling the load does not read fails CI instead of a multi-rank load. With language_config on the Omni spec, the test fails.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review (trigger comment carried no scoping instructions). 2 changed files, both reviewed in full: modelopt/torch/utils/moe.py (+140/-6) and tests/unit/torch/models/test_ep_specs.py (+235/-2). Read the base-branch EPSpec definition (modelopt/torch/models/specs.py:163-231), get_spec, the five registered EP sections (deepseek_v3, qwen3_moe, gpt_oss, llama4/llama4_text, nemotron_h_omni), and the one existing consumer of the refactored helpers (modelopt/torch/utils/plugins/model_load_utils.py:1044-1047).
Findings
- CRITICAL: 0
- IMPORTANT: 0
- SUGGESTION: 2 — a redundant function-level
from packaging.version import parse(already imported atmoe.py:32), and a missing data-level assertion that every registered spec'splan_config_attris a spelling the load actually reads.
What I verified
- The
_SUB_CONFIG_ATTRSextraction is behavior-preserving._ep_configand_num_experts_from_configboth readtext_configthenllm_config, same order as before, so the existingload_and_prepare_epcall sites are unaffected. The applier's refusal of any other spelling closes a silent-failure mode — a plan on the wrong config matches no parameter and the MoE is planned dense — which is the right trade. - Validate-before-mutate holds on every refusal path.
check_distributed_supportruns first, then theplan_config_attrspelling check, then target resolution, and only then the attach andprepare. Each of the three refusals (unsupported, version floor, unreadable/absent target) leaves the config untouched, and the tests assertvars(config) == {}rather than just the raise. _spec_refuses_ep'sexcept RuntimeErrordoes catch theunsupportedrefusal —NotImplementedErrorsubclassesRuntimeError. The inline comment saying so is worth keeping; it reads as a bug otherwise.- Existing-plan precedence is checked where the load reads a plan, not where the spec would write one.
getattr(_ep_config(hf_config), "base_model_ep_plan", None)covers the root and either nested config, so a stock plan onllm_configis not shadowed by ours ontext_config. This is the subtle case and it is both implemented and tested (on-another-nested-config). An explicitly-empty{}plan is treated as absent, consistently with_ep_config's own truthiness test. - The floor arithmetic.
max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse)compares parsed versions, not strings, so"5.18"correctly beats"5.7"— the string comparison would have inverted it. A model with no spec still gets the load's floor, viaget_specreturningNone→getattr(None, "ep_spec", None)→None. nemotron_h_omnihas norequires_transformersbutmin_transformers_version="5.18", so the floor stays at 5.7 for it. That is deliberate per the spec's comment (pre-5.18 the checkpoint loads through its own code, which names the configllm_config), and the dual-spellingplan_config_attrplus the("llm_config",)/("text_config",)/both parametrization cover it._EP_MIN_TRANSFORMERS = "5.7"is a new module constant that both the applier and the test read, so the skip marker and the enforced floor cannot drift.- The attached plan is copied (
dict(ep.base_model_ep_plan)), so two models prepared in one process cannot share a mutable dict and the registered spec cannot be mutated through a config — tested directly. - The user-facing error text points at flags that exist:
--use_fsdp2inexamples/hf_ptq/hf_ptq.pyandMODELOPT_FSDP_SIZEinmodelopt/torch/utils/env.py. prepareordering and config hand-back are tested, including that a returned replacement config propagates and thatpreparesees the plan already attached. TheEPSpec.preparedocstring already carries the obligation to forwardbase_model_ep_planwhen it returns a different config, so the one gap there is documented rather than silent.- Backward compatibility: purely additive public surface, no
modelopt_stateschema touched, no mode registration, no export path, no changed defaults._spec_refuses_epand_spec_ep_planhave no caller yet, which the description states and [3/3] supplies — acceptable for a declared split, not dead code.
Risk
Low. Nothing in the repository calls the new functions yet, the two refactored helpers keep their exact read order, and every refusal path is covered by a test with a stated mutation that breaks it. The one genuine residual risk is spec-data drift rather than logic: a future plan_config_attr typo is caught only when that model is loaded expert-parallel, which is what the test suggestion above addresses.
🤖 Generated with Claude Code
|
/ok to test 5648c93 |
5648c93 to
fe16583
Compare
|
/ok to test fe16583 |
69a4306 to
fe2e966
Compare
fe16583 to
1bd24c0
Compare
|
/ok to test 1bd24c0 |
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Comment: the module-wide Transformers skip removes existing dependency-free spec coverage, and the FSDP2 fallback advice misses a known refusal.
Needs action:
- Restore Transformers-free registry and plan checks in
tests/unit/torch/models/test_ep_specs.py; limit dependency skips to tests that actually need Transformers (inline). - Correct
check_distributed_support's FSDP2 advice to account for unsupported unquantized causal-LM checkpoints, including Nemotron-H MTP (inline). - Update the PR description's check order:
unsupportedis checked before the Transformers floor, not after it.
No action needed:
- Strengthening the existing plan assertions and extending them to
prepared_ep_planis justified; the blanket dependency skip is not.
| import pytest | ||
|
|
||
| from modelopt.torch.models import get_spec, get_specs | ||
| transformers = pytest.importorskip("transformers") |
There was a problem hiding this comment.
Bot comment.
This module-wide skip also suppresses all the pre-existing registry/plan tests when Transformers is absent. Those tests still need no Transformers: the model-spec package registers data independently of HF, and _routed_expert_keys is also dependency-free. noxfile.py::partial_unit(subset="torch") collects this directory without installing the HF extra, so this quietly removes coverage from that lane without a corresponding behavior change. Please keep the declarative checks runnable there and scope the dependency skip to the applier/version tests (or move those tests into a separate module).
There was a problem hiding this comment.
Agreed. The module-wide skip took the registry and plan checks out of a run without transformers, though they need none of it.
Fixed in bdd0448. transformers is now imported optionally, and only the tests that read the installed release skip without it (needs_transformers and needs_ep_transformers).
With transformers not importable, test_ep_specs.py now runs 30 tests and skips 14. Before this change it skipped the whole module.
| without_ep = ( | ||
| "without expert parallelism: in a single process, or FSDP2-sharded (--use_fsdp2 in " | ||
| "examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world) if the checkpoint is an " | ||
| "unquantized causal LM" |
There was a problem hiding this comment.
Bot comment.
Being an unquantized causal LM is not sufficient for the advertised FSDP2 fallback. For example, check_distributed_support("nemotron_h") below 5.7 emits this advice, but an unquantized Nemotron-H checkpoint containing MTP tensors then fails in examples/hf_ptq/example_utils.py::validate_fsdp2_supported, which explicitly refuses Nemotron-H MTP. Could the advice say “a checkpoint supported by the FSDP2 loader” or explicitly exclude MTP, rather than making this condition sound sufficient? The single-process alternative can remain.
There was a problem hiding this comment.
Agreed: being an unquantized causal LM is necessary, not sufficient.
Fixed in 3894689. The advice now offers FSDP2 "if the FSDP2 load supports the checkpoint, which takes an unquantized causal LM with no MTP head it cannot build (such as Nemotron-H's)". The single-process alternative still comes first.
#2746 [3/3] offers the same way out when an MoE would be left whole on every rank, and it now uses the same wording (d1dbf75).
|
/ok to test 3894689 |
|
/claude review |
|
Answers to the bot review of 1bd24c0, at 3894689:
Tests run (CPU, torch 2.14):
Before bdd0448, the first row skipped the whole module. Pre-commit, mypy included, passes against main. |
| if ep.base_model_ep_plan and not getattr(_ep_config(hf_config), "base_model_ep_plan", None): | ||
| target.base_model_ep_plan = dict(ep.base_model_ep_plan) | ||
|
|
||
| if ep.prepare is not None: | ||
| hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) | ||
| return hf_config |
There was a problem hiding this comment.
[SUGGESTION] The declarative path is guarded strictly against the silent-dense-MoE failure, but the prepare path is not — and it can reproduce the identical failure.
The unread / missing-target ValueErrors above exist because "attaching it elsewhere would be read as no plan at all and plan the MoE as dense." A prepare hook reaches the same end state with no error at all:
- it returns a different config object without carrying over the
base_model_ep_planjust attached totarget(exactly the obligationEPSpec.prepare's docstring places on the spec author) —_ep_config()on the returned config then finds no plan, and the MoE is planned dense; - it mutates in place and forgets to
return, soapply_distributed_supporthands the callerNone, which surfaces as an unrelatedAttributeErrordeep in the load.
Both are spec-authoring mistakes, which is precisely the class the unread check already refuses to let through silently. A post-condition costs two lines and keeps the two paths symmetric:
if ep.prepare is not None:
prepared = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code)
if prepared is None:
raise ValueError(f"{model_type}'s spec prepare returned no config to load.")
if _spec_ep_plan(model_type) and not getattr(
_ep_config(prepared), "base_model_ep_plan", None
):
raise ValueError(
f"{model_type}'s spec prepare returned a config carrying no expert-parallel plan "
"where the load reads one; a prepare returning another config must carry over the "
"base_model_ep_plan."
)
hf_config = prepared
return hf_configNon-blocking: no registered spec sets prepare today, and the PR already defers the other prepare-ordering question to #2748 — this would fit naturally there alongside it.
There was a problem hiding this comment.
Not a defect at this head, so I'm postponing it to this PR's next code change rather than pushing for it alone.
No spec reachable from this PR sets prepare. The first ones are Kimi-VL and Kimi-K2.5, in #2748. Their hooks return the config they were given, with the plan attached to its text_config by a nested apply_distributed_support("deepseek_v3", ...), and _ep_config reads the plan there. So no registered spec reaches either end state you describe: no plan where the load reads one, or a None config.
The post-condition guards the next spec author. I'll add it, with a test, in this PR's next code change.
| _KNOWN_STYLES = frozenset(moe._EXPERTS_STYLES) | {"grouped_gemm"} | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("model_type", EP_MODEL_TYPES) | ||
| def test_every_ep_plan_is_well_formed(model_type): | ||
| """Globs map to style names; a malformed entry matches nothing and fails silently at load.""" | ||
| plan = get_spec(model_type).ep_spec.base_model_ep_plan | ||
| if plan is None: | ||
| return | ||
| assert plan, f"{model_type} declares an empty plan, which is not the same as declaring none" | ||
| for glob, style in plan.items(): | ||
| assert isinstance(glob, str) and glob, f"{model_type}: bad glob {glob!r}" | ||
| assert isinstance(style, str) and style, f"{model_type}: bad style for {glob!r}" | ||
| """Every entry uses a style the expert-parallel load implements, and the plan names a | ||
| routed-experts container it can count and shard: a malformed plan fails silently at load.""" | ||
| ep = get_spec(model_type).ep_spec | ||
| for plan in (ep.base_model_ep_plan, ep.prepared_ep_plan): | ||
| if plan is None: | ||
| continue | ||
| assert plan, f"{model_type} declares an empty plan, which is not the same as declaring none" | ||
| for glob, style in plan.items(): | ||
| assert isinstance(glob, str) and glob, f"{model_type}: bad glob {glob!r}" | ||
| assert style in _KNOWN_STYLES, ( | ||
| f"{model_type}: {glob!r} uses style {style!r}, which the expert-parallel load does " | ||
| f"not implement (known: {sorted(_KNOWN_STYLES)})" | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] _KNOWN_STYLES is narrower than the styles a base_model_ep_plan may legitimately carry, and the failure message states something that isn't true.
moe._EXPERTS_STYLES is documented as "the styles a base_model_ep_plan gives a layer's routed-experts container" — it is deliberately the container vocabulary, not the set of styles the load accepts. The plan itself is consumed by transformers' distributed loading, which honours its own styles (local_colwise, rowwise, gather, …). So a plan entry that styles something other than a container — a router/gate entry, which EPSpec.base_model_ep_plan's own docstring treats as a real option ("A router entry is not always wanted", with DeepSeek-V3 given as the reason it is omitted there) — would fail this assertion saying the style is one "the expert-parallel load does not implement", when in fact the load never sees it as a container style and transformers implements it fine.
That makes the next spec author's failure message actively misleading. Two ways out, either is fine:
- keep the strictness but scope the claim to what it actually checks — assert the style set only over the entries
_routed_expert_keyspicks up, and word the message as "is not a routed-experts container style"; or - keep the whole-plan sweep and say instead that a style must be either a container style or one transformers implements, so a router entry passes.
Separately, "grouped_gemm" is spelled as a literal both here and in moe._routed_expert_keys; a shared module constant beside _EXPERTS_STYLES would keep the test and the implementation from drifting on it.
There was a problem hiding this comment.
Agreed: the message misleads for a router entry. The expert-parallel load keeps ep_router (it maps it to torch_a2a_router) and drops only the other styles outside the experts. A spec that added layers.*.mlp.gate: ep_router would load fine and still fail this test.
No registered spec has a router entry today; DeepSeek-V3 omits it on purpose. So every current plan passes, and nothing loads wrongly.
Postponed to this PR's next code change. I'll accept ep_router, reword the message to name a style the load does not keep, and spell grouped_gemm once, shared with _routed_expert_keys.
|
/ok to test 7c0be5b |
|
@coderabbitai Taken in 7c0be5b: the outside-diff item from your review on 41b8a6a,
The new test |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain
This addresses the reported concern. Following the repository’s review-feedback guidance, I checked the current source. I did not run the tests or verify the revert check. Per Review baseline: 🧠 Learnings used
You are interacting with an AI system. |
|
/claude review |
|
/ok to test 2b3dfc6 |
7c0be5b to
2b3dfc6
Compare
|
/claude review |
There was a problem hiding this comment.
Claude review passed — the two items open since the last review are fixed, and I found nothing new.
Scope: full review (trigger carried no scoping). 2 changed files (+431/-16), both read in full: modelopt/torch/utils/moe.py (+145/-6), tests/unit/torch/models/test_ep_specs.py (+286/-10). Also read EPSpec and all eight registered EP sections, _EXPERTS_STYLES / _routed_expert_keys, and _ep_config's fall-through, from the checkout. Baseline: .agents/review-guidelines/code-review.md at base e352caca.
No action needed:
- ✔️ Resolved since
38946893: the root-plan precedence gap (7c0be5b2) and thepreparereplacement-config assertion (41b8a6a1). - I traced the
existinghoist: with no plan anywhere the nested-target refusal is unchanged, and the attach is still guarded bynot existing, sotargetis only used after resolution.test_a_plan_at_the_root_needs_no_nested_configdoes fail without it — the oldnext(...)over(text_config, llm_config)returnedNonefor a root-only config. test_every_ep_plan_is_well_formedis strictly strengthened (style whitelist replacesisinstance(str), plus a routed-container check, extended toprepared_ep_plan); nothing was dropped.- Still open as non-blocking suggestions from the last round: no runtime check that
prepare's replacement config carries a load-readable plan (no spec setsprepareyet; author defers to #2748), and_KNOWN_STYLESis narrower than the plan vocabulary, so a future router style would fail the test — loudly, not silently. - Verification limits: I did not execute the tests (sandbox denied pytest) and did not check CI; both are the author's reported 5.18/5.5.4 runs.
|
/ok to test 289b626 |
2b3dfc6 to
289b626
Compare
kevalmorabia97
left a comment
kevalmorabia97
left a comment
There was a problem hiding this comment.
Three follow-ups you replied to claude[bot] with "postponed to this PR's next code change" are not in this head (it is only a rebase). Please land them here, or say explicitly which PR they move to (#2746 or #2748).
Needs action:
-
apply_distributed_support: afterprepareruns, check the returned config is notNoneand still carries a plan the load can read. -
test_ep_specs.py_KNOWN_STYLES: acceptep_routerand reword the failure message. -
_spec_ep_plandocstring: cover the case where the plan ships with transformers (qwen3_moe,minimax_m3_vl).
| target.base_model_ep_plan = dict(ep.base_model_ep_plan) | ||
|
|
||
| if ep.prepare is not None: | ||
| hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) |
There was a problem hiding this comment.
Please check here that the config prepare returns is not None and still carries a plan the load can read; today a bad return only fails later and further from the cause.
There was a problem hiding this comment.
Done in b29d579. apply_distributed_support now checks the config prepare returns, right where it returns it. It refuses with a ValueError naming the model type when that config is None, or carries no base_model_ep_plan where the load reads one (on the config, or its text_config or llm_config).
test_a_prepare_returning_no_config_the_load_can_read_is_refused covers both cases, and both fail without the check.
|
|
||
|
|
||
| # The styles the expert-parallel load implements for a plan's entries. | ||
| _KNOWN_STYLES = frozenset(moe._EXPERTS_STYLES) | {"grouped_gemm"} |
There was a problem hiding this comment.
Please accept ep_router here and reword the failure message at line 71-73 to match.
There was a problem hiding this comment.
Done in 202a805. _KNOWN_STYLES now includes ep_router, so it lists what the expert-parallel load keeps: the routed experts' container styles, grouped_gemm for their weights, and the router's ep_router. The failure message now says the load drops the style, and lists the ones it keeps.
| return hf_config | ||
|
|
||
|
|
||
| def _spec_ep_plan(model_type: str) -> dict[str, str]: |
There was a problem hiding this comment.
Please extend this docstring to cover the case where the plan ships with transformers (qwen3_moe, minimax_m3_vl).
There was a problem hiding this comment.
Done in ccab53a. The docstring now says the result is empty where the plan ships with transformers rather than the spec (qwen3_moe, minimax_m3_vl): the config carries the plan, and apply_distributed_support keeps it as it is.
|
All three follow-ups are now in this PR, at ccab53a. There is a reply with the commit in each thread. |
|
/ok to test ccab53a |
|
/claude review |
| ep = _ep_spec(model_type) | ||
| if ep is None or ep.unsupported: | ||
| return {} | ||
| return dict(ep.base_model_ep_plan or ep.prepared_ep_plan or {}) |
There was a problem hiding this comment.
[SUGGESTION] base_model_ep_plan or prepared_ep_plan reports only one of the two, but nothing enforces that a spec sets only one.
apply_distributed_support attaches base_model_ep_plan and then runs prepare, which attaches prepared_ep_plan — so for a spec that declares both, the plan the load ends up with is the union, while _spec_ep_plan returns the base plan alone. The consumer in [3/3] uses this to decide whether an MoE held whole on every rank should be refused, via _routed_expert_keys; if the extra containers only prepared_ep_plan names are the routed ones, that decision reads an incomplete plan and can rule the model dense.
EPSpec.prepared_ep_plan's docstring already says it is "for a model whose plan only prepare supplies", so the invariant is intended — it is just unchecked. Cheapest fix is a line in the registry test next to test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal:
assert not (ep.base_model_ep_plan and ep.prepared_ep_plan), (
f"{model_type} declares both plans; prepared_ep_plan is for a plan only prepare supplies"
)Non-blocking: no registered spec sets prepare today, so nothing hits this yet.
There was a problem hiding this comment.
Taken in f9c898a. test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal now also refuses a spec declaring both base_model_ep_plan and prepared_ep_plan. That is the invariant prepared_ep_plan's docstring states, so a load can't end up with the union while _spec_ep_plan reports only one. No registered spec, here or in the Kimi specs stacked above, declares both.
There was a problem hiding this comment.
LGTM — the three suggestions from the previous Claude round are resolved in code, and the only thing left is human sign-off on the strengthened existing test.
Needs action:
- Sign off on the edit to
test_every_ep_plan_is_well_formedintests/unit/torch/models/test_ep_specs.py— the diff justifies it (the looseisinstance(style, str)assertion became a style whitelist plus a routed-container check, extended toprepared_ep_plan); coverage only grows, but a bot cannot auto-approve a test edit. - Consider asserting
base_model_ep_planandprepared_ep_planare mutually exclusive in the registry test —_spec_ep_planreports only the first, whileapply_distributed_supportwould attach both (inline onmodelopt/torch/utils/moe.py:295, non-blocking).
No action needed:
- ✔️ Resolved since the last review: the unguarded
preparereturn value (now refusesNoneand a config with no plan where_ep_configreads one, tested both ways), the style whitelist and its message, and the_spec_ep_plandocstring's transformers-ships-the-plan case. The new root-plan shortcut is correct: theplan_config_attrspelling check still runs first, so validate-before-mutate holds on all five refusal paths. - Baseline:
.agents/review-guidelines/code-review.mdat basecf38ec3. Reviewed headccab53a, full scope, both changed files. I did not run tests or read CI; the test results in the description are unverified.
🤖 Generated with Claude Code
|
/ok to test f9c898a |
apply_distributed_support reads a model's EPSpec and prepares the config it is given for the expert-parallel load. It checks the transformers floor: the load's own 5.7, or the spec's later release, for a model with no spec too. It attaches the spec's base_model_ep_plan to the config the load reads it from, unless the model already carries a plan wherever the load reads one, and refuses a placement the load would not read. Then it runs the spec's prepare. check_distributed_support is the floor and refusal check alone, _spec_refuses_ep its predicate, and _spec_ep_plan the plan a spec would supply, read without running prepare. The nested configs the load reads a plan and an expert count from are one tuple, _SUB_CONFIG_ATTRS. Nothing calls them yet; the expert-parallel decision and load do in a later change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The well-formedness test only checked for non-empty strings: it now checks every entry's style against the ones the expert-parallel load implements, that the plan names a routed-experts container, and does so for the plan a spec's prepare attaches too. A new registry test checks that every spec puts its plan on a config the load reads. check_distributed_support drops its local import of parse, which the module imports at the top, and _ep_spec's local import says why it is local. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The applier's tests skipped the whole module when transformers was missing, which took the registry and plan checks, which need none of it, out of the partial-install unit run. Skip only the tests that read the installed transformers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The way around expert parallelism offered FSDP2 for any unquantized causal LM, but --use_fsdp2 refuses a Nemotron-H checkpoint carrying MTP weights. Name that as well. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test of the spec's prepare hook had it return a sentinel, so it held only that apply_distributed_support returns what prepare returned. A replacement config must also carry over the plan the spec attached, as EPSpec.prepare's contract asks: the stand-in now returns one that does, and the test checks the returned config is that one and holds the plan. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
apply_distributed_support required the nested config a spec names for its plan before it looked for a plan the config already carries, so a config with a plan of its own at its root, which wins and which the expert-parallel load reads there, was refused for lacking a nested config nothing needed. A plan the config carries is now found first; with none, a config without the nested config is still refused, naming both. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The load builds from the config a spec's prepare returns. One returning None, or a config carrying no base_model_ep_plan where the load reads one, failed only later in the load, far from the spec at fault. It is now refused as prepare returns, naming the model type. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The registered-plans check took only the routed experts' styles and grouped_gemm, though the expert-parallel load also keeps ep_router, which it swaps for its torch_a2a router. The check now accepts it, and its failure says the load drops the style and which it keeps. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_spec_ep_plan returns only the plan a spec supplies. For qwen3_moe and minimax_m3_vl transformers ships the plan on the config, which apply_distributed_support keeps, so their specs supply none. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
apply_distributed_support attaches base_model_ep_plan and then runs prepare, which attaches prepared_ep_plan, so a spec declaring both would load with their union while _spec_ep_plan reports one of them. The registry test now refuses such a spec, as prepared_ep_plan's docstring already intends it for a plan only prepare supplies. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 4cfbddf |
f9c898a to
4cfbddf
Compare
… refused Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> #3135 now refuses a spec prepare that returns no config, or one with no plan where the expert-parallel load reads one. Its refusal comes after the record of what the prepare patched for every later load: the patches stay in the process whatever it returned. A test has a Kimi-K3 prepare return None and checks that both types it spoils are recorded; it fails with the refusal above the record. The reload-refusal test's stand-in prepare now attaches the plan the spec states, as the real one does.
Note
Slice [2/3] of #2746's split (2026-10-30), so the model-spec work is reviewed and merged in pieces under the 500-line budget. Merge order: #2741 [1/3] (#3131, merged) → #2746 [1/3] (#3134) → this PR → #2746 [3/3] (#2746).
Base: #2746 [1/3] (#3134): the spec sections this PR reads. It also needs #2741 [1/3]'s
require_transformers, onmainsince #3131 merged. The diff below is this PR's alone.Series layout (the part around this PR)
main([OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [1/3], [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131, merged)What does this PR do?
Type of change: New feature
modelopt.torch.utils.moegains the one applier of the model specs' expert-parallel section (EPSpec, #2746 [1/3]), and the checks around it.apply_distributed_support(model_type, hf_config, *, ckpt_path=None, trust_remote_code=False)prepares a config for the expert-parallel load from the model's spec, in this order:unsupported. The refusal names a single process first, and FSDP2 (--use_fsdp2orMODELOPT_FSDP_SIZE=world) only where the FSDP2 load supports the checkpoint: an unquantized causal LM with no MTP head it cannot build. GPT-OSS ships MXFP4, which the FSDP2 load refuses, Llama-4's checkpoint is multimodal, and--use_fsdp2refuses Nemotron-H's MTP weights._EP_MIN_TRANSFORMERS), or the spec's later release. The load's floor holds for a model with no spec too. Its refusal offers the same ways around expert parallelism.prepare, after the plan is attached._ep_config): on the root, or on a nested config other than the one the spec would take. Overwriting it would substitute our idea of the layout for the model's.plan_config_attris neithertext_confignorllm_config, the two nested configs the expert-parallel load reads a plan and the expert count from (_SUB_CONFIG_ATTRS, which_ep_configand_num_experts_from_confignow read too). Attaching it elsewhere would be read as no plan at all and plan the MoE as dense.plan_config_attrmay list several spellings, the first one the config has taking the plan. A config with none of them is refused, naming them all.preparesupplies the plan, before anything is attached or prepared.check_distributed_supportis the floor and refusal check alone, for a config already prepared with its plan._spec_refuses_epis its predicate._spec_ep_planis the plan a spec would supply: its own, or the one it statesprepareattaches, read without runningprepare. A model below its spec's transformers floor keeps its plan; a spec refusing expert parallelism supplies none.Nothing calls these yet. #2746 [3/3] (#2746) has the expert-parallel decision and load call them.
Usage
Testing
tests/unit/torch/models/test_ep_specs.pygains the applier's tests:prepareordering;preparesupplies the plan;llm_configortext_config), and ontext_configfor a config with both;prepare.llm_configalone, or tryingllm_configbeforetext_config;EPSpec.prepareruns before the plan is known, so a load that ends up not expert-parallel keeps its effects. No spec here setsprepare. It is to be settled in [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748, which registers the first ones (Kimi-VL, Kimi-K2.5).tests/unit/torch/modelsandtests/unit/torch/utils, torch 2.14):prepared_ep_plantoo; a new test checks every spec puts its plan on a config the load reads (a bad style or an unread spelling fails them);check_distributed_supportdrops a redundant local import. Pre-commit passes against main;tests/unit/torch/{models,utils}pass on transformers 5.18 (777) and 5.5.4 (748, 29 skipped).tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (273) and 5.5.4 (255, 18 skipped).test_ep_specs.pyruns 30 tests and skips 14 with transformers not importable, where it skipped the whole module before.tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (273) and 5.5.4 (255, 18 skipped).tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (312) and 5.5.4 (282, 19 skipped).tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (312) and 5.5.4 (282, 19 skipped).tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (312) and 5.5.4 (282, 19 skipped).tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (314) and 5.5.4 (283, 20 skipped).preparehook now has it return a replacement config that carries the plan over, and checks that the returned config is that one and holds the plan. Pre-commit passes against main.tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (314) and 5.5.4 (283, 20 skipped).apply_distributed_supportfinds a plan the config already carries before requiring the nested config the spec names. A config with a plan of its own, at its root included, keeps it; with no plan anywhere, a config without the nested config is still refused.test_a_plan_at_the_root_needs_no_nested_configfails without the change. Pre-commit passes against main.tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (315) and 5.5.4 (283, 21 skipped).tests/unit/torch/models,test_moe_a2a.pyandtest_model_load_utils.pypass on transformers 5.18 (382) and 5.5.4 (339, 32 skipped).tests/unit/torch/models,test_moe_a2a.pyandtest_model_load_utils.pypass on transformers 5.18 (382) and 5.5.4 (339, 32 skipped).apply_distributed_supportrefuses apreparethat returnsNone, or a config with no plan where the load reads one, and names the model type (b29d579). The new test's two cases fail without the check.ep_router, and its failure names the styles the load keeps (202a805)._spec_ep_plan's docstring covers a plan that ships with transformers (ccab53a).tests/unit/torch/models,test_moe_a2a.pyandtest_model_load_utils.pypass on transformers 5.18 (384) and 5.5.4 (339, 34 skipped). The check changes nothing for [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746: its tests with the check applied pass on 5.18 (588).test_ep_specs.pypasses on transformers 5.18 (47) and 5.5.4 (32, 15 skipped).tests/unit/torch/models,test_moe_a2a.pyandtest_model_load_utils.pypass on transformers 5.18 (384) and 5.5.4 (339, 34 skipped).Before your PR is "Ready for review"
_ep_configand_num_experts_from_configread the same two nested configs, in the same order.CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models. Split out of #2746 (OMNIML-6076), whose review history stays on #2746.
Summary by CodeRabbit