Repository navigation
[OMNIML-6054] Decide how a model is distributed, from the model built on the meta device - #2737
shengliangxu wants to merge 36 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation |
|---|---|---|
| Linked Issues check |  |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
| Security Anti-Patterns | |
PASS. The PR changes only three modelopt Python modules and one test file. The changed code has no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), eval/exec on externa… |
| Title check | |
The title clearly describes the main change: determining model distribution from a model built on the meta device. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only three modelopt Python modules and one test file. The changed code has no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True),eval/execon external input, or# nosec. Transformers calls pass the caller-controlledtrust_remote_codevalue, and the public resolver defaults it toFalse; no hardcodedtrust_remote_code=Truewas added. Nopyproject.tomlorrequirements.txtdependency changes are present.
✨ Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- 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.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-1-session #2737 +/- ##
==============================================================
+ Coverage 77.61% 77.68% +0.07%
==============================================================
Files 661 662 +1
Lines 75257 75503 +246
==============================================================
+ Hits 58412 58656 +244
- Misses 16845 16847 +2
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:
|
f1e16c8 to
8b9651f
Compare
|
|
/ok to test 8b9651f |
shengliangxu
changed the title
2dcf64f to
a97136a
Compare
2ba6ef1 to
fe5b268
Compare
a97136a to
fc0ed3b
Compare
|
/ok to test fe5b268 |
shengliangxu
changed the title
fe5b268 to
74bbbf9
Compare
shengliangxu
changed the title
|
/ok to test 74bbbf9 |
74bbbf9 to
9074ecf
Compare
shengliangxu
changed the title
|
/ok to test 9074ecf |
c3f96b6 to
d023003
Compare
9074ecf to
00417b9
Compare
Two tests played rank 0 and rank 1 against a stand-in broadcast and asserted the message it was handed, and two more repeated broadcast_from's own failure tests. One pooled job replaces them: both gloo ranks resolve the plan of a real checkpoint, rank 1 with a knob of its own that would decide FSDP2, and both must load the plan rank 0 decided. The failure tests now live with broadcast_from's, on real ranks too, and the checkpoint-read job takes its bounded group from the shared helper. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…caller Nothing in the distribution decision calls moe.require_transformers or asks _auto_class for a VLM without its LM head: the expert-parallel loader is the first to, so both, with the test case for the headless backbone, arrive with it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Fold near-duplicates into one test per behavior, and drop tests that other tests already cover: - test_the_expert_count_is_read_under_every_spelling takes the dense, four-spelling and two VLM-nesting cases. - test_the_largest_ep_degree_dividing_the_world_and_the_experts_is_chosen takes the five degree cases. The ep * dp == world test goes: every case checks both degrees. - test_a_dense_model_is_fsdp_sharded_only_when_too_big_for_one_gpu takes the fits and too-big cases. - test_the_dense_knob_replicates_or_shards_over_the_whole_world takes the part-of-the-world refusal and an exported-empty case, and test_an_ep_degree_of_0_is_auto an exported-empty one. Both env-knob unit tests go: the decision tests fail if an empty or unset knob stops meaning auto, if a knob is read at import, or if one knob reads the other's variable. - test_a_model_without_an_ep_plan_has_no_routed_experts goes: test_an_moe_without_an_ep_plan_is_not_expert_parallel fails if an empty plan finds experts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ame a bad EP knob _auto_class took a VLM checkpoint's remote AutoModelForCausalLM whenever its auto_map named one, whether or not the load trusts the checkpoint's code, so an untrusted load built through a class it could not import instead of the native image-text class. It now takes the remote class only when trust_remote_code is set. A MODELOPT_EP_DEGREE that is not a number failed with int()'s bare ValueError, which the other ranks get as "invalid literal"; the error now names the knob. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
dense="auto" compared only what every rank holds whole with 60% of a GPU, leaving out the rank's 1/ep_size share of the routed experts, so it could replicate the dense part of an expert-dominated MoE that fits only with that part sharded: Qwen3-235B-A22B at EP 8 holds about 73 GB a rank on 80 GB GPUs with the dense part replicated, 59 GB with it sharded. The fit now counts that share. resolve_distribution_plan takes revision, token, cache_dir, local_files_only and subfolder, as load_pretrained does, and reads the checkpoint and its config with them, so a plan is decided for the checkpoint the load reads rather than the Hub default branch's, and a gated repo, another cache or a checkpoint in a subfolder can be planned. Its docstring no longer quotes a load_pretrained call this branch does not take. A single process parses its own ep_size through _ep_degree, so a bad one gets the error a distributed run gives it. The EP-degree helper returns the degree alone, the data-parallel degree being DistributionPlan.edp, and it and the expert-count helper are private like their neighbours. MODELOPT_FSDP_SIZE no longer accepts the undocumented "full". Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_auto_class took AutoModelForImageTextToText only for a config transformers maps to it natively, so a trusted remote VLM whose auto_map registers an image-text class, and whose config class is its own, fell through to AutoModel: the headless backbone, which leaves the LM head out of the dense fit. It now takes the registered class when the checkpoint's code is trusted, the class the remote causal-LM branch already defers to for such a checkpoint. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Building a remote-code model through an auto class registers its class there, as it is, and a later load through that auto class then builds the registered class, without the GenerationMixin a first load adds. Rank 0 builds one more meta model than its peers, to decide a distribution, so its model alone lacked generate. Under expert parallelism Kimi-K3's rank 0 then failed or skipped hf_ptq's preview while the other ranks decoded in lockstep, and the run deadlocked. build_meta_causal_lm now drops the registrations its build adds, so every rank's load builds the same class. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
build_meta_causal_lm read transformers' private _model_mapping._extra_content through getattr defaults, so a release that moved it would have turned the restore into a silent no-op, and with it brought back rank 0's model without generate and the deadlock it caused. It now reads the attribute directly: a move fails the build at once. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The grouped-GEMM parameters' parents counted only when no container style appeared in the plan, so a plan styling some layers' expert containers and only the grouped-GEMM parameters of others lost the second kind, both in the decision and in the expert-parallel load. Every transformers 5.18 plan with both styles names the same containers both ways, so the union changes nothing for them. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its *ForConditionalGeneration architecture reads like a VLM's, so T5, BART, Whisper and Qwen2-Audio fell through to AutoModel, which leaves the LM head out of the parameters the dense fit counts. The seq2seq auto classes now claim their configs first. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
MODELOPT_FSDP_SIZE=Auto or WORLD worked, but MODELOPT_EP_DEGREE=AUTO was refused as neither auto nor a degree. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
For a VLM whose nested text config holds the MoE, an expert-count field on the root config was read first, so an automatic ep_size could be chosen against a count the plan does not describe. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The refusal always said ep_size=, also for a degree read from MODELOPT_EP_DEGREE, so a user who set the variable was told to change an argument they never passed. It now names what the degree came from. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
It is the module's rank-0 printer, already used for the load's other notice, and flushes as the explicit print did. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The meta build kept only the auto class it built through. A remote VLM's __init__ can build its text model through AutoModelForCausalLM while the whole model builds through another auto class, and that registration stayed: rank 0's later load through AutoModelForCausalLM would then build the registered class, without the generate its peers' load adds. The build now also keeps AutoModel's and AutoModelForCausalLM's registrations as it found them. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A config with routed experts whose base_model_ep_plan names containers the built model has none of (a renamed submodule, a hand-written plan) left the decision with no experts to count, so the automatic degree became 1 and every expert was replicated on every rank without a word; the expert-parallel load finds the experts the same way, so it could not have sharded them either. That is refused now, naming the containers, unless the degree is set to 1. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
An explicit EP degree that does not divide the world size was refused by DistributionPlan, as ep_size=..., also when MODELOPT_EP_DEGREE set it. The decision now refuses it first and names the knob or argument the degree came from, as its refusal for the experts does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_auto_class fell back to the headless AutoModel for any multimodal config that neither the image-text nor a seq2seq class recognizes, also when AutoModelForCausalLM builds it whole, as it does Phi-4-multimodal (the one such config in transformers 5.5.4 and 5.18). The decision then weighed the model without its LM head. AutoModelForImageTextToText exists in every supported release, so its import guard goes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… decides the plan Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The largest degree holds the fewest experts per rank, leaving the most memory for activations, which calibration batches fill; any degree above 1 pays the same all-to-alls per MoE layer. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 1a4e341 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| from modelopt.torch.utils.moe import ( | ||
| _choose_ep_degree, | ||
| _ep_config, | ||
| _num_experts_from_config, | ||
| _routed_expert_modules, | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] These four MoE helpers are imported inside the function with no reason named, and _routed_expert_keys is imported a second time inside the refusal branch at line 1030. _auto_class does the same with AutoModelForImageTextToText / AutoModelForSeq2SeqLM / AutoModelForSpeechSeq2Seq, while AutoModel and AutoModelForCausalLM sit at module top in the same file.
CONTRIBUTING's coding standards ask for imports at module top, and for a brief comment naming the reason when one goes inside a function (circular import, optional dependency, unusually heavy import). None of these qualifies on its face: modelopt.torch.utils.moe does not import model_load_utils, so there is no cycle to break, and this module already hard-imports transformers at top (line 36), so the three extra auto classes are not an optional-dependency guard either — deferring them only moves an ImportError on an old transformers from load time to the middle of a plan decision, which is the failure mode the standard exists to avoid.
Worth either hoisting all of them to module top:
from transformers import (
AutoConfig,
AutoModel,
AutoModelForCausalLM,
AutoModelForImageTextToText,
AutoModelForSeq2SeqLM,
AutoModelForSpeechSeq2Seq,
PretrainedConfig,
)
from modelopt.torch.utils.moe import (
_choose_ep_degree,
_ep_config,
_num_experts_from_config,
_routed_expert_keys,
_routed_expert_modules,
)or, if one of them really does break a cycle or guard a version floor, adding the one-line comment that says which — the way from accelerate import init_empty_weights # only real callers of this function need it does at line 216. Either way the duplicate _routed_expert_keys import at line 1030 can fold into the single site.
| containers = ", ".join(sorted(_routed_expert_keys(ep_plan))) or "none" | ||
| raise ValueError( | ||
| f"{hf_config.model_type} has routed experts and a base_model_ep_plan, but none of the " | ||
| f"plan's expert containers ({containers}) names a module of the built model, so they " | ||
| f"can be neither counted nor sharded. Set {knob}=1 to load it without expert " | ||
| "parallelism." | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] This refusal conflates two causes, and the message only describes one of them.
_routed_expert_keys recognizes a container key only when its style is in _EXPERTS_STYLES = ("moe_tp_experts", "nemotron_experts_ep"), or when it is the parent of a grouped_gemm-styled key (moe.py:95-101). So a base_model_ep_plan whose experts are styled with anything else — stock megamoe_experts, which this repo's own register_torch_a2a docstring names as transformers' 5.10.4+ equivalent of torch_a2a_experts (moe.py:384-388), or packed_colwise/packed_rowwise expert weights with no grouped_gemm key, or a style a future release renames — yields an empty key set. The plan then lands here and reports:
… none of the plan's expert containers (none) names a module of the built model, so they can be neither counted nor sharded. Set MODELOPT_EP_DEGREE=1 to load it without expert parallelism.
which blames the model's module names when the actual cause is that ModelOpt recognized no container key in the plan at all. The remedy it offers (=1) silently gives up expert parallelism on a model whose plan is perfectly valid, and nothing points the reader at _EXPERTS_STYLES. (This is the loud-failure half of the earlier silent-ep_size=1 concern — good that it now raises; it is the wording and the single remedy that are off.)
Separating the two states costs about three lines and makes the diagnostic actionable, e.g.:
from modelopt.torch.utils.moe import _EXPERTS_STYLES, _routed_expert_keys
containers = sorted(_routed_expert_keys(ep_plan))
where = (
f"the plan's expert containers ({', '.join(containers)}) name no module of the built "
"model"
if containers
else f"the plan styles no routed-experts container ({', '.join(_EXPERTS_STYLES)} or a "
"grouped_gemm parameter's parent); its styles are "
f"{', '.join(sorted(set(ep_plan.values())))}"
)
raise ValueError(
f"{hf_config.model_type} has routed experts and a base_model_ep_plan, but {where}, so "
f"they can be neither counted nor sharded. Set {knob}=1 to load it without expert "
"parallelism."
)Since the expert-parallel loader in #2741 finds the experts through the same helper, whichever message fires is also the diagnostic for a load that would have replicated every expert — worth having it name the right knob to turn.
There was a problem hiding this comment.
Claude review — CRITICAL 0 · IMPORTANT 0 · SUGGESTION 2 (+1 carried over). Approving.
Scope: full review (the trigger comment carried no scoping instructions). 4 changed files (+1165/−8), all four read: modelopt/torch/utils/env.py, modelopt/torch/utils/moe.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py. Diff taken two-dot against origin/shengliangx/hfep-1-session. For composition context I read DistributionPlan / broadcast_from / all_reduce / set_distribution_plan / distribution_plan / size on the base branch, register_torch_a2a's pre-existing body in moe.py, modelopt/torch/utils/__init__.py, and CONTRIBUTING's coding standards. Nothing in the diff attempts to steer this review.
Prior rounds. I read the earlier Claude reviews. Since the last one the only code change is the _choose_ep_degree docstring clause, which answers that round's SUGGESTION 3. Both of its IMPORTANT findings now appear under Declined review findings in the PR description with substantive reasoning, and per CLAUDE.md that settles them — I re-checked both and accept the rationale rather than re-raising:
- Unconditional meta build. 0.54–0.65 s and 8–208 MB of host memory on 232B–1T configs is a real measurement against my estimate, the FSDP2 load builds the same meta model anyway, and
distribution=is a genuine route for a model that cannot be built on meta — which was the sharper half of my objection, and it has an answer. _auto_class'sAutoModelfallback for a remote*ForCausalLMVLM. Scoping it to the checkpoints the series actually supports (Kimi-VL/K2.5/K3 and Step-3.7-Flash all register*ForConditionalGeneration; DeepSeek-VL2 registers no auto class; Qwen3-Omni is unsupported on main) is a legitimate call, and the auto-class table now pins nine rows including the Phi-4-multimodal case.
What I checked and found sound
- Collective ordering — the part most likely to hang, and it is right.
all_reduce("checkpoint read", …)→broadcast_from("plan", …)is issued in the same order by every rank;ensure_local_checkpoint's own internal gloo group and broadcast are reached unconditionally by every rank becauseall_reducerunsfneverywhere, so a rank that cannot see the checkpoint fails after that inner collective and still joins the outer gather rather than leaving rank 0 blocked inbroadcast_object_list._refuse_a_launch_without_its_groupcovers the called-before-init_process_groupcase. - Rank-0-only inputs.
_decide_distributionreads only rank-local state (the knobs,_gpu_memory(), the built model) and runs insidebroadcast_from;world_size = dist.size()is read identically on every rank outside the closure. Divergent plans and the resulting mesh-mismatch hang are both closed, and the twocpu_dist_workerstests pin exactly those two properties. - Plan validity. Every exit from
_decide_distributionpasses throughDistributionPlan.__post_init__, soMODELOPT_EP_DEGREE=-2, anep_size>1in a one-rank run, and a caller-supplieddenseoutside{replicate, fsdp}all die loudly before any rank starts a load. Theif "auto"/elif world % ep/elif experts % epchain is correct in that order, andknobis captured beforeep_sizeis reassigned, so the refusals name the right source. - Dense fit.
model.parameters()dedupes tied embeddings and gets GQA right by construction rather than by estimate;id(p) in shardedkeys off the livenn.Parameterobjects theexpertsdict holds; the router (ep_router) and shared experts are correctly outside the sharded set; EP is resolved before the fit readsep_size, andep_size=1correctly makes the whole MoE the "dense" part. The dtype comes from the samebuild_meta_causal_lmthe FSDP2 load uses, so the fit and the load cannot disagree on bytes-per-element. - Suffix matching.
endswith("." + key)breaks on a dot, sosublayers.0.mlp.expertsis not matched bylayers.*.mlp.experts; the wildcard form is now computed once per module rather than once per key, with the same matches; and the base-relative keys match undermodel.,model.language_model.andlanguage_model.model.alike. - Registration restore.
_LazyAutoMapping.registeronly mutates_extra_content, so clear-and-restore is complete; snapshottingAutoModelandAutoModelForCausalLMalongsideauto_classcovers the towers a remote VLM's__init__builds through, the duplicate entry when they coincide is harmless, and thefinallyruns on a failed build too. Reading_extra_contentwithout a fallback is the right call for a private attribute whose silent absence would reintroduce the deadlock. - Backward compatibility and laziness.
auto_classis appended with a default that preservesbuild_meta_causal_lm's one in-tree caller;env.pyimports onlyosand is not star-imported intomodelopt/torch/utils/__init__.py, consistent withmoe.pyanddistributed.pyand keeping the genericep_degree/fsdp_sizenames out of the package namespace. No mode registration, nomodelopt_stateschema change, no export path touched — additive, as the description says. - Tests. Proportionate and CPU-only: tiny real configs built on meta, the process group's rank/size answered as torchrun would, two real gloo ranks only where the broadcast itself is the subject. The dense-fit tests derive the GPU size from the measured byte count instead of restating the formula, and each refusal is pinned by the knob or argument name it reports. Source insertions outside
tests/are 419, inside the ~500-line budget.
Findings
- [SUGGESTION] Function-level imports with no reason named (
model_load_utils.py:990). The four MoE helpers, a second_routed_expert_keysimport in the refusal branch, and threetransformersauto classes in_auto_classare all imported inside functions, whileAutoModel/AutoModelForCausalLMsit at module top in the same file. There is no cycle (moe.pydoes not importmodel_load_utils) and no optional dependency (transformersis a top-level import here), so CONTRIBUTING's import rule asks for either hoisting or a one-line reason comment. - [SUGGESTION] The ep-plan refusal conflates two causes and describes only one (
model_load_utils.py:1032-1038). A plan whose experts carry a style outside_EXPERTS_STYLESwith nogrouped_gemmkey — stockmegamoe_experts, which this repo's ownregister_torch_a2adocstring names, or packed expert weights — produces an empty key set and a message reading "expert containers (none) names a module of the built model", blaming the model's naming when ModelOpt recognized no container key at all, and offering=1(give up EP) as the only remedy. Three lines separate the two states and point at_EXPERTS_STYLES. - [SUGGESTION, carried over — no new comment posted] The expert-resolution guard is still one-sided (
model_load_utils.py:1025-1043). Count known, no container matched now refuses loudly, which is the good half of the earlier finding. Containers matched, count 0 is still unchecked:autoyieldsep_size=1silently, and an explicit degree raises"must divide the routed experts <model> shards (0)"immediately after_routed_expert_modulesfound them. It needs an expert-count field spelled outside the four names_num_experts_from_configknows, or a plan nested one level deeper than_ep_configdescends, so it stays a suggestion — butbool(experts) != bool(num_experts)makes both directions behave alike for one line.
Risk: low. The distributed core is the strong part and I could not fault it on this pass either: decide on rank 0, broadcast, every rank reads the checkpoint inside the failure-carrying reduce, every plan validated before any load starts. Nothing consumes the plan in this slice, which bounds the blast radius, and the one real bug it fixes — the auto-class registration leak that cost rank 0's Kimi-K3 its generate and deadlocked an EP preview — is fixed at its cause with two tests that fail without it. All three findings are non-blocking: one is a standards cleanup, two are diagnostics for cases a user would otherwise have to read _EXPERTS_STYLES to understand. Good to merge after #2736.
🤖 Generated with Claude Code
Note
The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [5/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.
Base: #2736. The diff below is this PR's alone, and it merges after #2736.
Re-cut: this PR's diff is unchanged; only its base moved. Review fixes since are listed at the end of Testing.
Series layout
mainWhat does this PR do?
Type of change: New feature
resolve_distribution_plan(ckpt_path)decides how a run spreads a model over its ranks. It returns aDistributionPlan(#2736) for every model, MoE or dense, and for a single process too. It is the one place a run decides this: #2779 loads the model by the plan and records the plan on it, and every later step reads it from there.auto_mapregistersAutoModelForImageTextToTextbuilds through that class, with its LM head. It decides from the model's parameters and then discards it, and leaves the auto class it built through as it was: building a remote-code model through an auto class registers its class there, as it is, and rank 0's later load through that class would then build it without theGenerationMixinevery other rank's first load adds. Rank 0's Kimi-K3 lackedgeneratewhile the other ranks' had it, and an expert-parallel preview deadlocked.ep_size, orMODELOPT_EP_DEGREE).autotakes the largest degree that divides both the world and the routed-expert count.base_model_ep_plannames. They are found by the rule the expert-parallel loader ([OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741) uses, so the decision and the load cannot disagree.dense, orMODELOPT_FSDP_SIZE, which takesauto,1to replicate, orworldor the world size to shard).autoFSDP2-shards the dense part when a whole copy of it plus the rank's share of the routed experts would take more than 60% of one GPU's memory.all_reduce([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), before the plan's broadcast. A rank that cannot see the checkpoint fails every rank with its own error, instead of leaving rank 0 waiting in the broadcast.resolve_distribution_plantakesrevision,token,cache_dir,local_files_onlyandsubfolderasload_pretraineddoes, so the plan and its config are for the checkpoint the load reads.broadcast_from, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736). The inputs are read locally: the knobs, this GPU's memory, the files on this node. Ranks deciding for themselves could disagree, and ranks loading different plans hang. If rank 0 fails, every rank raises.MODELOPT_EP_DEGREEorMODELOPT_FSDP_SIZEmust not reshape it. It holds the whole model, unless the caller's owndense="fsdp"asks for FSDP2 on one rank.Measuring the meta model avoids the gaps of estimating from config fields, which this PR's earlier revision did: it ignored GQA and tied embeddings, assumed bf16, and treated a config missing the fields it read as one that fits a GPU.
The class that builds a VLM with its head takes the checkpoint's remote class only under
trust_remote_code, aMODELOPT_EP_DEGREEthat is not a number is refused by name, and a single process parses its ownep_sizeas a distributed run does, refusing a bad one the same way.Usage
Testing
tests/unit/torch/utils/test_resolve_distribution_plan.py. It runs on CPU, builds tiny real configs on the meta device, and answers rank and size as torchrun would. It covers:cpu_dist_workers): both take rank 0's plan, though rank 1 exports a different dense layout, and rank 1 decides nothing; a rank that cannot see the checkpoint fails both ranks with its error, before the broadcast. Rank 0's own failures arebroadcast_from's, tested in [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736.generate, so rank 0 moved on to calibration while rank 1 decoded in lockstep, and the run deadlocked (TORCH_DISTRIBUTED_DEBUG=DETAIL: an all-reduce on one rank against an all-to-all on the other); with it, both ranks preview, calibrate and export.all_reduce, the experts' share left out of the dense fit, the Hub options not passed on, a single process'sep_sizenot parsed, and a trusted remote VLM's registered image-text class not taken.ep * dp == worldand no-EP-plan tests, which the decision tests cover, are dropped. Each drop is checked by mutation.MODELOPT_EP_DEGREEnames the knob, as a badep_sizein a single process names the argument.tests/uniton transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.MODELOPT_EP_DEGREEreads auto in any case; a hybrid plan names its routed experts by both styles; the expert count is read from the config holding the plan; a docstring note on a rank 0 without a GPU. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.MODELOPT_EP_DEGREE) or the argument the degree came from.AutoModel's andAutoModelForCausalLM's registrations as it found them; a routed-expert config whose expert-parallel plan names no module of the model is refused unless the degree is 1; an explicit degree that does not divide the world names the knob or argument it came from. Each new test fails without its change.build_meta_causal_lm's docstring states the registrations its restore assumes, andresolve_distribution_plan's that a caller passes the local checkpoint directory.AutoModelForCausalLMrecognizes (Phi-4-multimodal, the only such config in transformers 5.5.4 or 5.18) builds through it, head included, rather than throughAutoModel. The new table row fails without the change.resolve_distribution_plan's docstring says each rank returns its own copy of the config, rank 0's deciding the plan._choose_ep_degree's docstring says why the automatic EP degree is the largest: it holds the fewest experts per rank, leaving the most memory for activations, and every degree above 1 runs the same two all-to-alls per MoE layer.ep_size=1anddenseare both given. Rank 0 takes 0.54-0.65 s and 8-208 MB of host memory on 232B-1T configs, the FSDP2 load builds the same meta model anyway, and a model that cannot be built on meta loads withdistribution=, which skips the decision._auto_classfalls back toAutoModelfor a remote VLM whoseAutoModelForCausalLMclass is named*ForCausalLM. None of the remote checkpoints the series supports is one: Kimi-VL, Kimi-K2.5, Kimi-K3 and Step-3.7-Flash register*ForConditionalGeneration, and DeepSeek-VL2 registers no auto class. Qwen3-Omni, whichAutoModelcannot build, is not supported on main.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part 5 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2736.
Summary by CodeRabbit