Repository navigation
[OMNIML-6054] [5/21] Decide how a model is distributed, from the model built on the meta device - #2737
[OMNIML-6054] [5/21] Decide how a model is distributed, from the model built on the meta device#2737shengliangxu wants to merge 18 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 | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 71.70% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 53 functions across 4 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
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/execon externa…Title check The title clearly describes the main change: resolving model distribution from a model built on the meta device. It is specific and related to the pull request objectives. 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 ? 61.37%
=============================================================
Files ? 649
Lines ? 72149
Branches ? 0
=============================================================
Hits ? 44279
Misses ? 27870
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:
|
f1e16c8 to
8b9651f
Compare
|
|
/ok to test 8b9651f |
shengliangxu
changed the title
2ba6ef1 to
fe5b268
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 |
9074ecf to
00417b9
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — [5/21] distribution-plan decision (second pass, on 2aeb018c)
Scope: full review; the trigger carried no scoping instructions. All 4 changed files (+909/−7) reviewed: 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. Also read from the base branch, to check the contracts this diff leans on: DistributionPlan.__post_init__, broadcast_from, all_reduce, ensure_local_checkpoint, build_meta_causal_lm and load_pretrained.
Findings — CRITICAL: 0 · IMPORTANT: 2 · SUGGESTION: 3
| Severity | Where | What |
|---|---|---|
| IMPORTANT Performance | model_load_utils.py:691-701 |
dense="auto" drops the whole expert footprint from the budget, not the rank's 1/ep_size share, so auto can pick replicate for a plan that OOMs while fsdp fits |
| IMPORTANT Compatibility | model_load_utils.py:743-750 |
No revision/token/cache_dir/local_files_only/subfolder, so the documented pairing with load_pretrained can plan — and build the load's config — from a different revision than the weights |
| SUGGESTION | model_load_utils.py:714-718 |
The docstring's only snippet, load_pretrained(..., distribution=plan, hf_config=hf_config), is a TypeError at this commit; it lands in [6/21] |
| SUGGESTION | moe.py:75-89 |
choose_ep_dp_degrees' second return value is dead and duplicates DistributionPlan.edp; public-named but unexported (moe.py has no __all__) |
| SUGGESTION | model_load_utils.py:658-667 |
The single-process branch re-implements _ep_degree's parsing, losing the named error 2aeb018c just added; "full" is accepted but documented nowhere |
Most impactful
The dense-fit one. _REPLICATED_SHARE = 0.6 bounds only the dense part, never what the rank actually holds, and the excluded expert bytes are already in hand. Worked through on Qwen3-235B-A22B at 8×H100-80GB: experts ≈ 454 GB, dense ≈ 16 GB, ep_size="auto" → 8. The check asks 16 > 48? → no → replicate, and each rank then carries 454/8 + 16 = 72.8 GB — 91% of the GPU before a single activation. fsdp would have been 58.8 GB and fits. So auto picks the one layout that cannot run, on exactly the model class this stack targets. Adding sum(sharded bytes) // ep_size to the comparison is a three-line change, and I checked it against test_only_what_expert_parallelism_leaves_whole_must_fit: both of that test's assertions still hold, so it is additive rather than a semantics break.
The compatibility one is second because its revision case fails silently — the returned hf_config is documented as the config the load is built from, so planning on main while loading v2 decides the EP degree and the dense fit for the wrong model with no error anywhere.
Prior-round findings — the three raised last round are addressed
2aeb018c closed the one IMPORTANT from the earlier pass (_auto_class now gates remote_with_head on trust_remote_code, line 593-596) and the MODELOPT_EP_DEGREE naming suggestion (line 643-644), both with tests (test_the_class_that_builds_the_whole_model, test_an_ep_degree_that_is_not_a_number_names_the_knob). Two earlier items are still open and not re-filed as inline comments, since they were already put to you once:
num_experts = num_experts_from_config(hf_config) if experts else 0(line 683) still collapses "not MoE" with "the EP plan matched no module", so a config spelling its expert count outside the four attributesnum_experts_from_configknows silently drops toep_size=1rather than failing. Declining it is a reasonable call — flagging only that it is unchanged.- CodeRabbit's open item on
_auto_classreturningAutoModelfor a trusted remote VLM that advertisesAutoModelForImageTextToTextinauto_map(line 597-605), which undercounts the LM head in the dense decision. It overlaps the first IMPORTANT above, since both are about the budget being too small.
What checked out
Spelling these out because they are the places this kind of code usually breaks, and they are all handled:
- Collective safety of the new
all_reduce.checkpoint_readruns unconditionally asfnon every rank, soensure_local_checkpoint's internalnew_group→broadcast_object_list→destroy_process_group(hf_checkpoint_utils.py:922-936) is issued by all ranks in the same order. Both failure modes it can raise — "which this rank cannot see" (line 942) and a per-rankAutoConfigfailure — fire after those internal collectives have completed, so the raising rank still joinsall_gather_objectand no rank is stranded. This is the fix CodeRabbit asked for in the first round, and it is correct. - Plan validation. Every invalid combination I tried is rejected loudly by
DistributionPlan.__post_init__rather than silently carried:ep_size=2in a single process (1 % 2 != 0),ep_size < 1,dense="shard".MODELOPT_FSDP_SIZE=4at world 8 is refused by_dense_from_env. The world-1 branch's reliance onDistributionPlan'sworld_size=1default is right. - Argument order through the changed signature.
build_meta_causal_lm(ckpt_path, trust_remote_code, None, hf_config, auto_class=...)matches(ckpt_path, trust_remote_code, attn_implementation, hf_config=None, auto_class=...);dtypeis derived inside fromhf_config(line 201), not a parameter, so there is no positional shift.auto_classis additive and the pre-existing call at line 489 keeps its meaning. - Byte counting.
model.parameters()de-duplicates, so tied embeddings are counted once and not double-counted against the dense budget;numel()/element_size()are both valid on meta tensors;dtypefalls back tobfloat16rather than float32, so there is no 2× overestimate for a config without a dtype. - Expert matching.
_routed_expert_modulesmatches only the container (layers.*.mlp.experts), not its children, andm.parameters()then recurses — so sibling routers andshared_expertscorrectly stay in the replicated count.replace_layer_number_by_wildcardandbase_model_ep_planare both transformers-5 symbols andpyproject.toml:95pinstransformers>=5.5,<5.19, so there is no version-guard gap. - Knob laziness.
env.py's accessors reados.environper call, matching the docstring and the launcher-forks-workers case;env.pyis deliberately not star-imported, so the genericep_degree/fsdp_sizenames cannot collide inmodelopt.torch.utils. - Tests. One test per behaviour, CPU-only, real tiny configs on meta, with the two genuinely-distributed cases on the session-scoped
cpu_dist_workers— consistent with the repo's "keeptests/unitfast" and "be stingy" rules. The rank-1-exports-a-different-knob and rank-1-cannot-see-the-checkpoint cases are the right two to spend real ranks on.
Risk assessment: moderate. The structure is sound — one decision point, rank 0 decides, failures propagate, no rank-dependent collectives — and the collective ordering, the plan validation and the byte accounting are all correct. The risk is concentrated in the dense="auto" heuristic: it is the one path that can hand back a plan no rank can execute, and it does so on large MoEs rather than on an edge case. Both IMPORTANT findings are contained fixes to _decide_distribution / resolve_distribution_plan with no API break; neither requires reworking the design.
🤖 Generated with Claude Code
|
/ok to test 6aaf081 |
|
On CodeRabbit's outside-diff comment in review 5554774496 ( |
|
/claude review |
| def _ep_degree(ep_size) -> int | str: | ||
| """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``.""" | ||
| value = env.ep_degree() if ep_size is None else ep_size | ||
| if str(value) in ("auto", "0"): | ||
| return "auto" | ||
| try: | ||
| return int(value) | ||
| except ValueError: | ||
| name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" |
There was a problem hiding this comment.
[SUGGESTION] The two knobs normalize their value differently, so they accept different spellings of the same intent.
_dense_from_env (line 625) does env.fsdp_size().lower(), so MODELOPT_FSDP_SIZE=WORLD works — test_the_dense_knob_replicates_or_shards_over_the_whole_world pins that with the "WORLD" case. _ep_degree compares str(value) raw, so MODELOPT_EP_DEGREE=AUTO falls through to int("AUTO") and is refused as "expected auto, 0 (auto) or an EP degree." The asymmetry runs the other way for whitespace: int(" 4 ") succeeds, while MODELOPT_FSDP_SIZE=" 1 " is refused because .lower() does not strip.
Both knobs are set by hand in launcher shells, which is exactly where AUTO and a stray space come from, and the failure is a hard raise at plan resolution rather than a fallback. Normalizing once in each accessor makes the two consistent:
def ep_degree() -> str:
"""Expert-parallel degree knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``."""
return os.environ.get("MODELOPT_EP_DEGREE", "").strip().lower() or "auto"and the same .strip().lower() in fsdp_size(), letting _dense_from_env drop its own .lower(). _ep_degree then still needs str(value).strip().lower() for the ep_size argument path, which does not go through env.
| ep_plan = getattr(_ep_config(hf_config), "base_model_ep_plan", None) or {} | ||
| experts = _routed_expert_modules(model, ep_plan) | ||
| num_experts = _num_experts_from_config(hf_config) if experts else 0 |
There was a problem hiding this comment.
[SUGGESTION] The EP plan and the expert count are read from two different configs, so a VLM can take its plan from the nested text config and its count from the root.
Line 685 resolves the plan through _ep_config(hf_config), which deliberately walks into .text_config / .llm_config to find the submodel that carries base_model_ep_plan. Line 687 then calls _num_experts_from_config(hf_config) on the root config, and that helper tries the root's four attribute spellings first, only recursing if none hit. For every VLM shape in the tests the root carries no expert count, so the two agree — but they are not structurally tied, and a root config that exposes an expert count of its own (a composite config forwarding attributes, or an MoE vision tower) would have the EP degree chosen from one submodel's count and applied to another's experts. ep_size then divides a number unrelated to the modules _routed_expert_modules found.
Since _ep_config has already picked the config the decision is about, counting from it removes the possibility:
ep_config = _ep_config(hf_config)
ep_plan = getattr(ep_config, "base_model_ep_plan", None) or {}
experts = _routed_expert_modules(model, ep_plan)
num_experts = _num_experts_from_config(ep_config) if experts else 0_num_experts_from_config still recurses, so the plain-MoE case is unchanged and test_the_expert_count_is_read_under_every_spelling keeps passing. This is distinct from the "no experts vs. plan matched nothing" point from the earlier round — it is about the two helpers reading the same config, not about the attribute list.
| if dense == "auto": | ||
| # A rank holds a whole copy of what expert parallelism does not shard, and its 1/ep_size | ||
| # share of the routed experts it does. | ||
| sharded = ( | ||
| {id(p) for m in experts.values() for p in m.parameters()} if ep_size > 1 else set() | ||
| ) | ||
| held = sum( | ||
| p.numel() * p.element_size() // (ep_size if id(p) in sharded else 1) | ||
| for p in model.parameters() | ||
| ) | ||
| memory = _gpu_memory() | ||
| too_big = memory is not None and held > _REPLICATED_SHARE * memory | ||
| dense = "fsdp" if too_big else "replicate" |
There was a problem hiding this comment.
[SUGGESTION] Now that the experts' share is inside the budget, too_big no longer tells apart "the dense copy is what overflows" (FSDP2 fixes it) from "the rank's expert share alone overflows" (FSDP2 cannot).
This is the right fix for the gap raised last round — the comparison is now against what a rank actually holds. The side effect is that the two cases collapse. On an expert-dominated MoE, held - dense_bytes can already exceed _REPLICATED_SHARE * memory by itself; the branch then picks "fsdp", which shards a dense part that was never the problem, buying almost nothing while adding an all-gather per dense layer to every forward, and the run still OOMs with no indication why. test_what_a_rank_holds_with_the_dense_part_replicated_must_fit's the-experts-share-does-not case is exactly this shape: the experts' quarter is what overflows, and the assertion is "fsdp".
"fsdp" is still the best of the two available layouts, so the decision needn't change — but the user is the only one who can act on it (a larger ep_size is already maxed, so it means more GPUs). A warning when the unshardable part alone exceeds the budget would say so, and costs a few lines on a path that runs once:
memory = _gpu_memory()
budget = None if memory is None else _REPLICATED_SHARE * memory
dense = "replicate" if budget is None or held <= budget else "fsdp"
if budget is not None and dense == "fsdp":
unshardable = held - sum(
p.numel() * p.element_size() for p in model.parameters() if id(p) not in sharded
)
if unshardable > budget:
warnings.warn(
f"At ep_size={ep_size} a rank's share of the routed experts is "
f"{unshardable / 1e9:.1f} GB, over the {budget / 1e9:.1f} GB budget on its own, "
"so FSDP2-sharding the dense part will not make this model fit. More ranks, or "
"a larger ep_size, would."
)Both assertions of that test still hold either way — it only reads plan.dense.
There was a problem hiding this comment.
Claude review — [5/21] distribution-plan decision (third pass, on 6aaf081d)
Scope: full review; the trigger carried no scoping instructions. All 4 changed files (+978/−7) reviewed: 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. Also read from the base branch, to check the contracts this diff leans on: DistributionPlan.__post_init__, broadcast_from, all_reduce, ensure_local_checkpoint, build_meta_causal_lm, load_pretrained and modelopt/torch/utils/__init__.py.
Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3
| Severity | Where | What |
|---|---|---|
| SUGGESTION | model_load_utils.py:638-646 |
_ep_degree normalizes its value differently from _dense_from_env: MODELOPT_FSDP_SIZE=WORLD works but MODELOPT_EP_DEGREE=AUTO raises, and whitespace tolerance runs the other way |
| SUGGESTION | model_load_utils.py:685-687 |
The EP plan comes from _ep_config(hf_config) but the expert count from the root hf_config, so the two are not structurally tied to the same submodel |
| SUGGESTION | model_load_utils.py:695-707 |
With the experts share now inside the budget, too_big no longer distinguishes "the dense copy overflows" (FSDP2 helps) from "the ranks expert share alone overflows" (it cannot) |
None of these blocks merge.
Both IMPORTANT findings from the previous round are fixed, and I verified each
- The dense fit now counts the ranks expert share (line 697-703).
helddivides a sharded parameters bytes byep_sizeand leaves everything else whole, which is what the_REPLICATED_SHAREcomment and theresolve_distribution_plandocstring now both describe. Re-ran my Qwen3-235B-A22B arithmetic against the new code: 8×H100-80GB, experts ≈ 454 GB, dense ≈ 16 GB,ep_size="auto"→ 8, soheld = 454/8 + 16 = 72.8 GBagainst a0.6 × 80 = 48 GBbudget →"fsdp". The layout that cannot run is no longer the oneautopicks.test_what_a_rank_holds_with_the_dense_part_replicated_must_fitpins both directions by constructing the GPU size from the real meta models expert bytes rather than a magic number, and itsep_size=1assertion pins that without EP the whole MoE must fit. - The Hub options are now plumbed through, and they match
load_pretrainedline for line:ensure_local_checkpoint(revision, token, cache_dir, local_files_only), thenos.path.join(local_path, subfolder), then anAutoConfigread taking everything butrevision— the same split, for the same documented reason (anauto_mapnaming code in another repo).test_the_hub_options_locate_the_checkpoint_the_load_readsis a good test for this: a real two-revision offline cache with a subfolder, asserting onhidden_size, so planning against the wrong revision fails it rather than passing silently.
The three earlier SUGGESTIONs that were taken are also in: _ep_degree now names MODELOPT_EP_DEGREE vs ep_size, the single-process branch routes through _ep_degree instead of re-parsing, choose_ep_dp_degrees dead second return is gone (_choose_ep_degree), and the docstring no longer carries the distribution=plan snippet that was a TypeError at this commit. Two remain open by your call and I have not re-filed them: the num_experts = ... if experts else 0 collapse, and _auto_class as a third parallel auto-class resolver in the repo.
What checked out
- Collective safety.
checkpoint_readruns unconditionally asall_reducesfn, soensure_local_checkpoints internalnew_group→broadcast_object_list→destroy_process_group(hf_checkpoint_utils.py:922-934) is issued by every rank in the same order. I checked the raise sites: both the "which this rank cannot see" path (line 941) and a per-rankAutoConfigfailure fire after those internal collectives complete, so the raising rank still joinsall_gather_objectand strands nobody.readis a per-call closure list with exactly one local entry, soread[0]is this ranks own read in both the distributed and the single-process (combine([fn()])) paths. - Validation precedes the expensive work and the mutation.
_ep_degreeand_dense_from_envboth raise before the meta model is built, and every plan that could not be executed is rejected byDistributionPlan.__post_init__rather than carried:ep_sizenot dividing the world,ep_size < 1(including a negative that slips past the> 1divisibility check),ep_size != 1on a single process, and adensestring that is neither layout — which is what makes the un-translateddense="world"argument a loud failure rather than a silent one. Partial-world FSDP is refused in_dense_from_env. Nothing is recorded on a model here, so a refusal leaves no state behind. - No config mutation leaking across ranks. The returned
hf_configis the object_decide_distributionhands tobuild_meta_causal_lmon rank 0 only, so iffrom_configmutated it in place, rank 0s config would diverge from the other ranks and[6/21]would build mismatched models. Transformers_from_configdeep-copies the config for exactly this reason, andattn_implementation=Nonehere meansbuild_meta_causal_lms ownhf_config._attn_implementationwrite (line 198-200) is skipped. No divergence. - Byte accounting.
model.parameters()de-duplicates, so tied embeddings count once;id(p)identity againstshardedis sound whilemodelis alive;m.parameters()recurses from the container, so sibling routers andshared_expertsstay in the replicated count while the experts own grouped-GEMM parameters leave it;numel()/element_size()are valid on meta tensors;dtypefalls back tobfloat16, not float32, so there is no 2× overestimate. - Expert matching.
_routed_expert_modulesmatches the container and not its children (layers.*.mlp.experts.gate_up_projwildcards to something that does not end in.layers.*.mlp.experts), and thegrouped_gemmfallback reconstructs the same container key.replace_layer_number_by_wildcardandbase_model_ep_planare both transformers-5 symbols against the>=5.5,<5.19pin, and theneeds_qwen3_moe_ep_planguard covers the 5.7 plan arrival. - Namespace and laziness.
modelopt/torch/utils/__init__.pydoes not star-importenv, so the genericep_degree/fsdp_sizenames cannot collide inmodelopt.torch.utils;env.pys accessors reados.environper call;transformers.integrations.tensor_parallelis imported inside_routed_expert_modulesandAutoModelForImageTextToTextbehindtry/except ImportError. The per-rankAutoConfigread withtrust_remote_codematches whatload_pretrainedalready does (line 433), so it introduces no new concurrency surface. - Tests. One test per behaviour, CPU-only, real tiny transformers configs built on meta rather than mocks, with only the two cases that genuinely need real ranks on the session-scoped
cpu_dist_workers— the rank-1-exports-a-different-knob case and the rank-1-cannot-see-the-checkpoint case, which are the right two. The_knobscontext manager clearing a pooled workers inherited environment, and the_no_distribution_knobsautouse fixture, are the kind of thing that is usually missing and causes flakes here.
Risk: low. The two paths that could hand back an unexecutable plan — the dense fit and the checkpoint the plan is decided from — are both fixed this round, and the structure around them (one decision point, rank 0 decides, failures propagate before any broadcast, validation in __post_init__) is sound. Still purely additive: resolve_distribution_plan, _auto_class, env.py and the new moe.py helpers have no caller outside the new tests until [6/21], and the only change to existing code is build_meta_causal_lms defaulted auto_class, which leaves every existing call sites meaning intact. No modelopt_state, mode registration, config schema or export path is touched, so there is no restore or checkpoint-compatibility surface. ~321 insertions of source, well inside the ~500-line budget.
I ran nothing — no GPU here, and _gpu_memory is the one input the CPU tests necessarily stub. @shengliangxu, the parts only real hardware can confirm are unchanged from last round: the _REPLICATED_SHARE = 0.6 headroom against actual calibration activations, and the total-vs-free-memory question Codex raised, now documented in the docstring as assuming the runs GPUs are alike.
🤖 Generated with Claude Code
|
/ok to test 03fcbe5 |
|
/ok to test 9e9643a |
|
/ok to test b0c59a4 |
decide_distribution_plan turns a model's config and the world size into a DistributionPlan. Each degree is an explicit value or "auto", and explicit always wins. - Capability: a model is expert-parallel capable when transformers ships a base_model_ep_plan for it, found on the root config or a VLM's nested text config, and it has experts under any of the spellings configs use. - EP degree: it must divide both the world and the expert count, so every rank holds whole experts; auto takes the largest such degree. - Dense: replicated when it fits one GPU with headroom, else FSDP2-sharded across the world. The decision runs on rank 0 alone and is broadcast; the other ranks adopt it. Its inputs are read locally -- the knobs, this GPU's memory, the config on this rank's filesystem -- so ranks deciding independently could disagree, and a disagreement does not raise: they build meshes of different shapes and hang at mesh creation or the first collective. If rank 0 fails to decide, it still reaches the broadcast and sends the failure, so every rank raises instead of waiting. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
decide_distribution_plan hand-rolled the rank-0 decision: broadcast the plan, or rank 0's error so the other ranks raise instead of hanging. Inside a ModelOptSession it now goes through session.decide, which does the same and keeps the result as session.plan, the same on every rank. Without a session the process decides alone, as a single process always did. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
decide_distribution_plan sets session.plan with session.plan.decide(fn): rank 0 decides, and every rank holds and returns rank 0's plan, read back as session.plan.value. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… it on Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…meta model resolve_distribution_plan is now the one place a run decides how a model is spread over its ranks, for every model rather than only an MoE. It moves here from load_pretrained's routing, with the MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE knobs, and decide_distribution_plan, a second public way to decide, is gone. Rank 0 decides from the model built on the meta device, by the class that builds it whole: a causal LM, or a VLM through the class that carries its LM head. The routed experts are the ones its base_model_ep_plan names, found the way the expert-parallel load will find them. Whether the dense part fits one GPU is judged from the parameter bytes that every rank would hold whole, rather than from a config estimate that ignored GQA and tied embeddings and gave up on configs without the fields it read. An explicit EP degree must divide the routed-expert count, and MODELOPT_FSDP_SIZE either replicates the dense part (1) or shards it over the whole world (world): a part of the world is refused. A single process decides without building anything: it holds the whole model, unless dense='fsdp' asks for FSDP2 on one rank. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The module docstring lists its knobs as an indented block under a paragraph line, which reStructuredText reads as a definition list. A knob wrapped onto a deeper-indented second line is unexpected indentation there, and the docs build fails on it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
MODELOPT_FSDP_SIZE FSDP2-shards the dense part for "world" or for the world size as a number, which is what launchers set (64 on 64 ranks), but the docs named only "world". Say so in the env module and in resolve_distribution_plan. The module lists its knobs as a bullet list now, so an entry can wrap without breaking the Sphinx build. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The plan decision sizes dense="auto" against rank 0's GPU, by its total memory, and every rank takes that decision; the docstring said only "one GPU's memory", leaving the assumption that the run's GPUs are alike unstated. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…s broadcast resolve_distribution_plan made the checkpoint local and read its config on each rank before the broadcasts rank 0 sends the plan through. A rank that could not see the checkpoint raised there alone, and rank 0 waited in the broadcast for it. The read now goes through the failure-carrying all_reduce, so every rank raises with the failing rank's error first. The tests that answer a launch's broadcasts answer its gather too. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A single process read MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE before returning its whole-model plan, so a shell still exporting a multi-GPU run's knobs reshaped a one-process load: an EP degree of 4 made a plan the world of one cannot hold and raised, MODELOPT_FSDP_SIZE=world asked a plain load for FSDP2, and a size of 2 raised. With nothing to shard over, the knobs are now not read; the caller's own arguments still count, so dense="fsdp" is still FSDP2 on one rank. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
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>
|
/ok to test d756401 |
What 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([4/21]) for every model, MoE or dense, and for a single process too. It is the one place a run decides this: [6/21] 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.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 ([9/21]) 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([4/21]), 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, [4/21]). 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 [4/21].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/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 5 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [4/21].
Summary by CodeRabbit