Skip to content

[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
shengliangx/hfep-1-sessionfrom
shengliangx/hfep-2-plan-decision
Open

shengliangxu wants to merge 14 commits into
shengliangx/hfep-1-sessionfrom
shengliangx/hfep-2-plan-decision

Conversation

@shengliangxu

@shengliangxu shengliangxu commented Oct 12, 2026 •

Copy link
Copy Markdown
Collaborator

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 a DistributionPlan ([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.

  • From the model, not estimates. Rank 0 builds the model on the meta device, using the class that loads it whole: a causal LM, or for a VLM the class that carries its LM head. It decides from the model's parameters and then discards it.
  • EP degree (ep_size, or MODELOPT_EP_DEGREE).
    • auto takes the largest degree that divides both the world and the routed-expert count.
    • The routed experts are the ones the config's base_model_ep_plan names. They are found by the rule the expert-parallel loader ([9/21]) uses, so the decision and the load cannot disagree.
    • An explicit degree must divide the routed-expert count.
  • Dense layout (dense, or MODELOPT_FSDP_SIZE, where 1 replicates and world shards).
    • auto FSDP2-shards what every rank would hold whole when that would take more than 60% of one GPU's memory.
    • Experts that expert parallelism shards don't count toward this; without expert parallelism, the whole model does.
    • Sharding over part of the world is refused.
  • Every rank reads the checkpoint first. Each rank makes the checkpoint local and reads its config inside the failure-carrying 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.
  • Rank 0 decides, and every rank gets its plan (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.
  • A single process decides without building anything, and reads no knob: a shell still exporting a multi-GPU run's MODELOPT_EP_DEGREE or MODELOPT_FSDP_SIZE must not reshape it. It holds the whole model, unless the caller's own dense="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

# torchrun --nproc_per_node 8 ...
from modelopt.torch.utils.plugins.model_load_utils import resolve_distribution_plan

plan, hf_config = resolve_distribution_plan("Qwen/Qwen3-30B-A3B")
# DistributionPlan(world_size=8, ep_size=8, dense='replicate')

Testing

  • New 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:
    • the knobs, read at call time; expert counting; EP-degree selection;
    • routed experts found under any model prefix, named by their style or as their grouped-GEMM parameters' parent;
    • the class that builds a whole model, a VLM's with its LM head;
    • the decision:
      • a single process, whatever knobs its shell exports; an EP-capable MoE; an MoE without an EP plan;
      • the dense fit from exact parameter bytes, counting only what expert parallelism leaves whole;
      • knob and argument overrides;
      • refused partial sharding and indivisible EP degrees;
      • a VLM decided from its whole model;
      • a model that cannot be built on meta, named in the error;
    • on two real gloo ranks (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 are broadcast_from's, tested in [4/21].
  • Mutation-checked: 30 mutants, each failing a test; and, since, a rank deciding for itself, the plan sent from the wrong rank, a single process reading the knobs, and the checkpoint read outside the all_reduce.
  • Near-duplicate tests are merged into one test per behavior, and the env-knob, ep * dp == world and no-EP-plan tests, which the decision tests cover, are dropped. Each drop is checked by mutation.
  • Locally on this revision, the reordered series rebased on main 264fd26: the full unit suite on transformers 5.18, and on 5.5 apart from the tests whose ONNX, TensorRT or diffusers packages that environment lacks; CI's transformers-free unit lane; pre-commit over the whole PR against main; and the Sphinx docs build. The same at every PR of the stack.

Before your PR is "Ready for review"

  • Is this change backward compatible?: ✅ Additive.
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
  • Did you write any new necessary tests?: ✅
  • Did you update Changelog?: N/A. The feature's entry lands with [21/21].
  • Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 5 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [4/21].

Summary by CodeRabbit

  • New Features
    • Added automatic model distribution planning, including expert- and data-parallel degree selection and dense-parameter sharding based on GPU memory.
    • Added support for environment-configured parallelism settings and model selection and sizing for vision-language models.
    • Added validation for model configurations and parallelism choices.
  • Bug Fixes
    • Distribution plans are now consistent across ranks, and planning failures are propagated to all ranks.

@copy-pr-bot

copy-pr-bot Bot commented Oct 12, 2026

Copy link
Copy Markdown

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.

@coderabbitai

coderabbitai Bot commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

The change adds model and routed-expert discovery, selects EP and dense-layout settings, and resolves a distribution plan across ranks. Unit tests cover discovery, planning decisions, model-class selection, broadcasts, and checkpoint visibility errors.

Changes: Model distribution planning

Layer / File(s) Summary
Model and routed-expert discovery
modelopt/torch/utils/moe.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py
Adds expert-count and EP-degree helpers, matches routed-expert keys to model modules, and selects the Transformers model class for meta-model construction. Tests cover config discovery, routed-expert matching, and model-class selection.
EP and dense-layout decisions
modelopt/torch/utils/env.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py
Adds environment accessors and decision logic for EP degree and dense replication or FSDP sharding. Tests cover defaults, overrides, divisibility, memory thresholds, and invalid settings.
Plan resolution and rank coordination
modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py
Adds checkpoint config loading and rank-0 plan computation and broadcast. Tests cover broadcast behavior and checkpoint-visibility errors.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant resolve_distribution_plan
  participant Rank0
  participant MetaModel
  participant OtherRanks
  Caller->>resolve_distribution_plan: checkpoint path and planning options
  resolve_distribution_plan->>Rank0: config and planning request
  Rank0->>MetaModel: build meta model and discover routed experts
  Rank0->>resolve_distribution_plan: distribution plan or error
  resolve_distribution_plan->>OtherRanks: broadcast plan
  OtherRanks->>resolve_distribution_plan: adopt plan
Loading
Merge Risk: 🟡 Moderate · up to 7c4fe

For some composite or remote vision-language checkpoints, the distribution plan can pick a wrong expert-parallel degree or a wrong dense layout. Fix both selection paths before merging.

Pre-merge checks: Passed 5 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning 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.
Passed Passed checks: 5 passed
Check name Status Explanation
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns Passed 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…
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check Passed The title clearly summarizes the main change: determining model distribution from a model built on the meta device. It is specific and related to the added distribution-planning functionality.
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/exec on external input, or # nosec. Transformers calls pass the caller-controlled trust_remote_code value, and the public resolver defaults it to False; no hardcoded trust_remote_code=True was added. No pyproject.toml or requirements.txt dependency 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.

@shengliangxu
shengliangxu added this pull request to stack #2755 October 12, 2026 00:57
@codecov

codecov Bot commented Oct 12, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.19626% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.31%. Comparing base (a43e6a3) to head (9bd8b2a).

Files with missing lines Patch % Lines
modelopt/torch/utils/moe.py 93.93% 2 Missing ⚠️
modelopt/torch/utils/plugins/model_load_utils.py 98.52% 1 Missing ⚠️
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     
Flag Coverage Δ
unit 61.31% <97.19%> (+0.13%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from f1e16c8 to 8b9651f Compare October 12, 2026 02:24
@github-actions

github-actions Bot commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2737/

Built to branch gh-pages at 2026-10-22 06:54 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 8b9651f

@shengliangxu shengliangxu changed the title [2/19] Decide the expert-parallel plan from the model and the knobs [OMNIML-6054] [2/19] Decide the expert-parallel plan from the model and the knobs Oct 12, 2026
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch 2 times, most recently from 2ba6ef1 to fe5b268 Compare October 13, 2026 01:32
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test fe5b268

@shengliangxu
shengliangxu removed this pull request from stack #2755 October 13, 2026 17:42
@shengliangxu
shengliangxu added this pull request to stack #2775 October 13, 2026 17:43
@shengliangxu shengliangxu changed the title [OMNIML-6054] [2/19] Decide the expert-parallel plan from the model and the knobs [OMNIML-6054] [2/20] Decide the expert-parallel plan from the model and the knobs Oct 13, 2026
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from fe5b268 to 74bbbf9 Compare October 13, 2026 19:38
@shengliangxu
shengliangxu removed this pull request from stack #2775 October 13, 2026 19:38
@shengliangxu
shengliangxu added this pull request to stack #2781 October 13, 2026 19:40
@shengliangxu shengliangxu changed the title [OMNIML-6054] [2/20] Decide the expert-parallel plan from the model and the knobs [OMNIML-6054] [2/21] Decide the expert-parallel plan from the model and the knobs Oct 13, 2026
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 74bbbf9

@shengliangxu
shengliangxu removed this pull request from stack #2781 October 14, 2026 00:13
@shengliangxu
shengliangxu changed the base branch from shengliangx/hfep-1-session to shengliangx/hfep-1b-session-facts October 14, 2026 00:13
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 74bbbf9 to 9074ecf Compare October 14, 2026 00:14
@shengliangxu
shengliangxu added this pull request to stack #2786 October 14, 2026 00:14
@shengliangxu shengliangxu changed the title [OMNIML-6054] [2/21] Decide the expert-parallel plan from the model and the knobs [OMNIML-6054] [3/22] Decide the expert-parallel plan from the model and the knobs Oct 14, 2026
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 9074ecf

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 9074ecf to 00417b9 Compare October 14, 2026 01:47
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 7c4fe81 to 587a194 Compare October 20, 2026 00:58
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 587a194

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 587a194 to 3b64221 Compare October 20, 2026 01:12
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 3b64221

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 3b64221 to 4f07f65 Compare October 20, 2026 01:54
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 4f07f65

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 4f07f65 to 42ae17a Compare October 20, 2026 21:09
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 42ae17a

@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-2-plan-decision branch from 42ae17a to 6f91f3a Compare October 21, 2026 17:43
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 6f91f3a

@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 78d7b9f

@shengliangxu

Copy link
Copy Markdown
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>
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 9bd8b2a

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant