Repository navigation
[OMNIML-6054] [5/21] Decide how a model is distributed, from the model built on the meta device - #2737
Open
shengliangxu wants to merge 14 commits into
Open
[OMNIML-6054] [5/21] Decide how a model is distributed, from the model built on the meta device#2737shengliangxu wants to merge 14 commits into
shengliangxu wants to merge 14 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. |
Contributor
shengliangxu
added this pull request to stack #2755
October 12, 2026 00:57
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-1-session #2737 +/- ##
==============================================================
+ Coverage 61.18% 61.31% +0.13%
==============================================================
Files 647 648 +1
Lines 71562 71773 +211
==============================================================
+ Hits 43783 44008 +225
+ Misses 27779 27765 -14
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:
|
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 12, 2026 02:24
f1e16c8 to
8b9651f
Compare
Contributor
|
Collaborator
Author
|
/ok to test 8b9651f |
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
2 times, most recently
from
October 13, 2026 01:32
2ba6ef1 to
fe5b268
Compare
Collaborator
Author
|
/ok to test fe5b268 |
shengliangxu
removed this pull request from stack #2755
October 13, 2026 17:42
shengliangxu
added this pull request to stack #2775
October 13, 2026 17:43
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 13, 2026 19:38
fe5b268 to
74bbbf9
Compare
shengliangxu
removed this pull request from stack #2775
October 13, 2026 19:38
shengliangxu
added this pull request to stack #2781
October 13, 2026 19:40
Collaborator
Author
|
/ok to test 74bbbf9 |
shengliangxu
removed this pull request from stack #2781
October 14, 2026 00:13
shengliangxu
changed the base branch from
shengliangx/hfep-1-session
to
shengliangx/hfep-1b-session-facts
October 14, 2026 00:13
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 14, 2026 00:14
74bbbf9 to
9074ecf
Compare
shengliangxu
added this pull request to stack #2786
October 14, 2026 00:14
Collaborator
Author
|
/ok to test 9074ecf |
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 14, 2026 01:47
9074ecf to
00417b9
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 20, 2026 00:58
7c4fe81 to
587a194
Compare
Collaborator
Author
|
/ok to test 587a194 |
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 20, 2026 01:12
587a194 to
3b64221
Compare
Collaborator
Author
|
/ok to test 3b64221 |
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 20, 2026 01:54
3b64221 to
4f07f65
Compare
Collaborator
Author
|
/ok to test 4f07f65 |
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 20, 2026 21:09
4f07f65 to
42ae17a
Compare
Collaborator
Author
|
/ok to test 42ae17a |
shengliangxu
force-pushed
the
shengliangx/hfep-2-plan-decision
branch
from
October 21, 2026 17:43
42ae17a to
6f91f3a
Compare
Collaborator
Author
|
/ok to test 6f91f3a |
Collaborator
Author
|
/ok to test 78d7b9f |
Collaborator
Author
|
/ok to test 5ca9d16 |
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>
Collaborator
Author
|
/ok to test 9bd8b2a |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.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, where1replicates andworldshards).autoFSDP2-shards what every rank would hold whole when that 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.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.
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.ep * dp == worldand no-EP-plan tests, which the decision tests cover, are dropped. Each drop is checked by mutation.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