[shengliangxu]

What does this PR do?

Type of change: New feature

load_pretrained (#2733) loads a model whole. A distributed load, such as parallel_load_and_prepare_fsdp2 (main), is a separate entry point, and every caller picks one itself, as hf_ptq.py --use_fsdp2 does. A caller that picks a loader also gets what that loader skips: a sharded load cannot run the model type's loading hook, so a Nemotron-H checkpoint's MTP head is not built and its tensors go to the export unquantized.

This PR makes load_pretrained the one loader, whole or distributed:

• distribution is a DistributionPlan; by default it is the plan resolve_distribution_plan ([5/21]) decides.
  • dense="fsdp": parallel_load_and_prepare_fsdp2 over the process group. A single process launched with torchrun --nproc_per_node=1 has no group, so the load sets one up, of one rank, once nothing can refuse the load anymore; a process started without a launcher cannot set one up and is refused. distributed_run ([4/21]) destroys it with the run. One rank of a multi-rank launch whose group is not set up yet is refused too, told to set it up first with distributed_run(), rather than shard over the launch's ranks under a one-rank plan; the plan-size refusal says the same when that is why the sizes differ.
  • Otherwise from_pretrained. On a distributed run every rank gets the whole model, on its own GPU.
  • A plan that shards the routed experts raises here. [12/21] adds that load, once the expert-parallel loader ([9/21]-[11/21]) is in. Until it lands, the changelog entry says so: an MoE planned expert-parallel, as the default MODELOPT_EP_DEGREE=auto plans one whose config has a base_model_ep_plan, is refused before it loads; it loads with MODELOPT_EP_DEGREE=1 (with MODELOPT_FSDP_SIZE=world to FSDP2-shard it) or a plan with ep_size=1.
  • resolve_distribution_plan takes the class the load builds with (auto_class, which load_pretrained passes). With a class the sharded loads do not build and no knob asking to shard, it plans the whole model on every rank without a meta-device build, so an encoder-decoder loaded through AutoModel, for example, loads multi-rank as before.
• The plan is recorded on the model and every submodule (set_distribution_plan, [4/21]), where calibration and export read it. load_pretrained records it exactly once, at the end of every load; the loaders it calls record none. fsdp2_wrap records nothing, as on main, so a caller that wraps or loads a model through it directly records the plan itself, as the error for a missing plan says. hf_ptq.py --use_fsdp2 loads through parallel_load_and_prepare_fsdp2 directly and records no plan in this PR; nothing on that path reads the plan before [8/21], which records it in the same PR that makes get_max_batch_size read it.
• fsdp2_wrap moves to model_load_utils, with the load that calls it and every other loader, and imports LayerActivationCollector at module level. modelopt.torch.utils.distributed.fsdp2_wrap forwards to it, so callers of the old path keep working, given the hf extra the loaders need, and the FSDP2 helpers it uses stay in distributed.
• Every load makes a Hub ID local and records the unplaced source keys ([OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733). The FSDP2 load builds its model from the config without the Hub options, so custom code an auto_map names comes with the environment's Hub settings.
• Every rank loads the same plan. Every rank takes rank 0's distribution, a plan or none, through one broadcast (broadcast_from, [4/21]) before it branches, and loads a plan rank 0 was handed just as it would one it decided. Ranks that each loaded the plan they were handed could build mismatched meshes and hang, and ranks passing different kinds would wait in each other's collectives. Rank 0 checks what it was handed before it sends it, a non-plan included, so a distribution it refuses raises on every rank instead of leaving the others waiting in the broadcast.
• Refusals, before anything loads.
  • A sharded load cannot run the model type's loading hook, so it refuses a checkpoint the hook would build modules for: the MTP head above.
  • It refuses from_pretrained's placement arguments (device_map, max_memory, ...) and a model class it does not build, rather than ignore them, and an explicit dtype, saying why it does not apply: the load keeps the checkpoint's dtype, and FSDP2's mp_policy sets the dtype it computes in.
  • The FSDP2 load refuses a quantized checkpoint, whose tensors it would read as stored.
  • The FSDP2 load refuses a VLM, which it would build through AutoModelForCausalLM: refused deep in the loader, or its language model alone. It refuses any config _auto_class does not build through AutoModelForCausalLM, before it sets up a group or reads anything, and the refusal says how to load the model whole on every rank: its model class as auto_class, with MODELOPT_FSDP_SIZE=1 or a replicated plan.
  • A plan made for another world size is refused, and so is anything other than a plan.
  • hf_config without the plan it was resolved with is refused. A sharded load builds the model from it as given, so a plan decided here, from the checkpoint's own config, would not be the one it came with. config= is for a config the plan should be decided from. resolve_distribution_plan's docstring says to hand the config it returns to load_pretrained only with a plan that shards the model: a whole-model load reads the checkpoint's config itself and refuses hf_config.

The FSDP2 route's refusal names its ways out (MODELOPT_FSDP_SIZE=1, or a replicated plan), and a config handed with a plan is copied before the load builds from it. With a plan handed and no config, the FSDP2 route reads the checkpoint's config once and hands it to the load, rather than reading it twice. Deciding a plan builds the meta-device model from a copy of the config, so the config returned with the plan, which a sharded load builds from, is the same on every rank: the build writes its dtype and its attention and experts implementations, and only rank 0 builds.

Usage

# torchrun --nproc_per_node 8 quantize.py
from modelopt.torch.utils.distributed import distributed_run
from modelopt.torch.utils.plugins.model_load_utils import load_pretrained

with distributed_run():
    model = load_pretrained("meta-llama/Llama-3.1-8B")  # a whole copy on every rank's GPU
    # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead.

Testing

• New tests/unit/torch/utils/test_load_pretrained_distribution.py. It runs on CPU, and stand-in loaders record their calls. It covers:
  • each load: the arguments it gets and the plan it records;
  • a decided FSDP2 load, handed the config the plan was decided from instead of reading it again, and a handed FSDP2 plan's load reading the checkpoint's config once (that check fails without its fix);
  • a single process's FSDP2 load setting up its group of one, and leaving none behind when the load is refused;
  • every refusal, raised before any loader runs, the expert-parallel plan's included; test_a_load_refuses_what_it_cannot_load has a VLM row (a Llava config), which fails without the refusal;
  • on two real gloo ranks (cpu_dist_workers), rank 0 handed a replicated plan and rank 1 an FSDP2 one, or none: both load whole and record rank 0's plan. A rank that loads the plan it was handed fails at once, on the placement refusal. Rank 0 handed a non-plan fails both ranks;
  • a caller's own model class: loaded whole on every rank, planned without building the model, with no sharding knob, and refused by the sharded load with one.
• tests/unit/torch/utils/test_resolve_distribution_plan.py: the config returned with the plan is the one every rank read, unchanged by rank 0's meta-device build (new test_the_config_returned_with_the_plan_is_as_every_rank_read_it, which fails without its fix).
• A single process's load, placed as from_pretrained places it, with the caller's arguments and a single process's plan: [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's real CPU load in tests/unit/torch/utils/test_unplaced_source_keys.py, which now also checks the plan.
• fsdp2_wrap leaving the plan to the load: a CPU unit test in tests/unit/torch/utils/test_load_pretrained_distribution.py. [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's GPU tests that wrap or load through it directly (test_export_carry_over.py, test_model_load_utils.py) record the plan themselves. Its old path forwarding every argument: tests/unit/torch/utils/test_distributed_run.py.
• Mutation-checked: 17 mutants, each failing a test. They drop the rank-GPU default, pin a single process's model, drop the config handoff or the world-size check, load an expert-parallel or FSDP2 plan whole, accept a string, let each rank load the plan it was handed, record no plan in the load or record one in fsdp2_wrap too, set no group up for a single process's FSDP2 load, set one up over a live group, or have the old fsdp2_wrap path drop an argument or the result.
• New tests/gpu/torch/utils/test_load_pretrained_distribution.py runs load_pretrained left to decide, on every GPU of the node:
  • a dense model loads whole on each rank, and with MODELOPT_FSDP_SIZE=world FSDP2-sharded;
  • each rank's weights match the checkpoint;
  • the load records its unplaced keys and its plan.
  • Not run locally; CI's GPU workflow runs it.
• tests/examples/hf_ptq/test_example_utils.py: get_model's stand-in model is a module now, since load_pretrained records its plan on it.
• The load refusals are one parametrized test, and the single-process load test gives way to [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's real load. Each drop is checked by mutation.
• A single process without a launcher, or one rank of a multi-rank launch whose group is not set up, is refused an FSDP2 load, without stubbing the setup, and a handed config is left as the caller passed it. The refusal table has an explicit dtype.
• Locally on this revision, the stack rebased on main 0638a27: pre-commit over the whole PR against main, at every PR of the stack; each fix's own tests at its PR; and at the stack's tip the unit tests of the changed areas (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"

• Is this change backward compatible?: ✅ load_pretrained keeps its behavior in a single process unless the plan shards the model (FSDP2 there needs torchrun --nproc_per_node=1), and distribution is new. fsdp2_wrap behaves as on main and records no plan; its old import path still works, given the hf extra.
• 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?: ✅ 0.49.0, Misc.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 6 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [5/21]. The expert-parallel half of this routing is [12/21].

Replaces #2774, which carried this routing whole as [19/20]. Moving its branch down the stack made the branch part of #2774's old base, so GitHub closed #2774 as merged into that base; nothing reached main.

Summary by CodeRabbit

• New Features

  • Added GGML imatrix calibration with weighted search and per-expert statistics for fused MoE models.
  • Updated IQ recipes for combined GPTQ and imatrix calibration, and added a Nemotron 3.5 Super VL AutoQuantize recipe.
  • Added standalone draft model training and export for external-draft speculative decoding.
  • Distributed loading supports replicated or FSDP2-sharded models through configurable distribution plans, which are recorded on loaded models. Automatic plan selection and Hub authentication and caching options are supported.

• Documentation

  • Added guidance for distributed loading and distribution-plan recording.

[copy-pr-bot]

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]

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.

• Last reviewed commit: fb8bf38
• Latest commit, not reviewed: 6f38823

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

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-13b-load-routing
• File: .coderabbit.yaml
• Commit: fb8bf38
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 45c74fdb-8dbc-4269-9909-f3d32a1010a4

📥 Commits

Reviewing files that changed from the base of the PR and between 6150cee and fb8bf38.

📒 Files selected for processing: 4

• CHANGELOG.rst
• modelopt/torch/utils/plugins/model_load_utils.py
• tests/unit/torch/utils/test_load_pretrained_distribution.py
• tests/unit/torch/utils/test_unplaced_source_keys.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

load_pretrained now accepts or resolves a distribution plan and selects whole-model or FSDP2-sharded loading. It validates inputs, records load metadata, and provides FSDP2 wrapping through the model-loading plugin.

Changes: Distributed model loading

Layer / File(s)	Summary
Plan-based load routing modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_load_pretrained_distribution.py, tests/gpu/torch/utils/test_load_pretrained_distribution.py, tests/examples/hf_ptq/test_example_utils.py, CHANGELOG.rst, tests/unit/torch/utils/test_unplaced_source_keys.py	load_pretrained resolves or accepts a distribution plan and selects whole-model or FSDP2-sharded loading. Sharded loads validate supported arguments and checkpoints. The loader records the plan and unplaced keys. Unit and GPU tests cover routing, validation, and placement.
FSDP2 wrapping and delegation modelopt/torch/utils/plugins/model_load_utils.py, modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_run.py, tests/unit/torch/utils/test_load_pretrained_distribution.py	The plugin implements FSDP2 wrapping, and distributed.py delegates to it. Tests check argument forwarding and confirm that wrapping alone does not record a distribution plan.
Distribution metadata in distributed tests tests/gpu/torch/export/test_export_carry_over.py, tests/gpu/torch/utils/test_model_load_utils.py	Distributed test models record an FSDP distribution plan before export-related checks.

Priority: Normal

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

Change: Feature · Severity of issue fixed: Medium

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant load_pretrained
  participant _route
  participant _load_sharded
  participant parallel_load_and_prepare_fsdp2
  participant from_pretrained
  Caller->>load_pretrained: Provide checkpoint and optional distribution plan
  load_pretrained->>_route: Validate plan and select load path
  alt Dense FSDP plan
    load_pretrained->>_load_sharded: Load checkpoint with config
    _load_sharded->>parallel_load_and_prepare_fsdp2: Load on local CUDA device
  else Whole-model plan
    load_pretrained->>from_pretrained: Load model with device map
  end
  load_pretrained-->>Caller: Return model with plan and unplaced keys recorded

Loading

Suggested reviewers: kevalmorabia97, sugunav14

Merge Risk: Minimal · up to fb8bf

The documented distributed loading paths are ready for normal checks. Users selecting an unsupported expert-parallel plan must choose a supported loading option.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 68.42% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 57 functions across 10 files. (1 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Title check		The title names the main change: load_pretrained now picks its loading route (whole-model or FSDP2) from a DistributionPlan. The changeset matches this. The title does not mention the moved fsdp2_wrap…
Linked Issues check		The only linked issue, #2774, is closed and completed. The issue states it is historical context only. No active directly linked target remains, so no linked-issue coding requirements apply to this pu…
Out of Scope Changes check		The whole-PR summary (mergeBase 4e10244 to currentHead fb8bf38) shows changes that match the PR description. In modelopt/torch/utils/plugins/model_load_utils.py, load_pretrained gains a distribution…
Security Anti-Patterns		No listed security anti-pattern is introduced by the pull request. In the changed modelopt Python files, torch.load uses weights_only=True; no added weights_only=False, allow_pickle=True, ha…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Title check

Explanation

The title names the main change: load_pretrained now picks its loading route (whole-model or FSDP2) from a DistributionPlan. The changeset matches this. The title does not mention the moved fsdp2_wrap or the new Hub options, but it does not need to cover every detail. The bracketed prefix [6/21] is a series marker and does not affect the assessment.

Full details: Linked Issues check

Explanation

The only linked issue, #2774, is closed and completed. The issue states it is historical context only. No active directly linked target remains, so no linked-issue coding requirements apply to this pull request. The prior PASS result still holds.

Full details: Out of Scope Changes check

Explanation

The whole-PR summary (mergeBase 4e10244 to currentHead fb8bf38) shows changes that match the PR description. In modelopt/torch/utils/plugins/model_load_utils.py, load_pretrained gains a distribution parameter, resolve_distribution_plan gains config and Hub access options, and fsdp2_wrap moves here. In modelopt/torch/utils/distributed.py, fsdp2_wrap forwards to the new location and its signature is unchanged. The Hub access options support the PR's Hub ID localization. The test changes cover the same functions: test_example_utils.py (FakeModel now a Module so the plan can be recorded), test_export_carry_over.py and test_model_load_utils.py (manual plan recording after fsdp2_wrap), test_distributed_run.py (forwarding), test_unplaced_source_keys.py (plan check), test_resolve_distribution_plan.py (Hub options), and the new unit and GPU test modules. The CHANGELOG.rst entry documents the same behavior. No file in the summary is shown to be unrelated. The assessment rests on the AI-generated summary; the source hunks were not read in this review.

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern is introduced by the pull request. In the changed modelopt Python files, torch.load uses weights_only=True; no added weights_only=False, allow_pickle=True, hardcoded trust_remote_code=True, eval()/exec() call, or # nosec comment exists. The pull request changes no pyproject.toml or requirements*.txt file, so it adds no dependency requiring license review.

Full details: Docstring Coverage

Explanation

Docstring coverage is 68.42% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 57 functions across 10 files. (1 skipped: 1 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

✨ Finishing Touches: 💡 2

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 95.49550% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.42%. Comparing base (0a1b3c8) to head (6f38823).

Files with missing lines	Patch %	Lines
modelopt/torch/utils/plugins/model_load_utils.py	95.37%	5 Missing ⚠️

Additional details and impacted files

@@                         Coverage Diff                          @@
##           shengliangx/hfep-2-plan-decision    #2779      +/-   ##
====================================================================
+ Coverage                             61.37%   61.42%   +0.05%     
====================================================================
  Files                                   649      649              
  Lines                                 72156    72237      +81     
====================================================================
+ Hits                                  44286    44374      +88     
+ Misses                                27870    27863       -7     

Flag	Coverage Δ
unit	61.42% <95.49%> (+0.05%)	⬆️

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]

/ok to test 3ac7c7b

[shengliangxu]

/ok to test df76b36

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2779/
Built to branch gh-pages at 2026-10-25 18:50 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[shengliangxu]

/ok to test aa3615e

[shengliangxu]

/ok to test 168c90d

[shengliangxu]

/ok to test 537ff8d

[shengliangxu]

/ok to test aa2a065

[shengliangxu]

/ok to test 6f38823