[shengliangxu]

What does this PR do?

Type of change: New feature

load_pretrained (#2733) loads a model whole. The distributed loads, load_and_prepare_ep ([6/20]-[8/20]) and parallel_load_and_prepare_fsdp2 (main), are separate entry points, and every caller picked one itself, hf_ptq.py with its own copy of the plan decision. A caller that picks a loader also gets what that loader skips. On the expert-parallel path that was the Nemotron-H MTP head: prepare_model_for_loading builds it only on a whole model, so the expert-parallel load carried it into the export unquantized, without a word.

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

• distribution picks the load.
  • "sequential": from_pretrained, as before.
  • "fsdp2": parallel_load_and_prepare_fsdp2.
  • A DistributionPlan: the experts sharded at load (load_and_prepare_ep, the dense part replicated or FSDP2-sharded as fsdp_size says), FSDP2 over the world, or a whole copy on every rank's own GPU.
• Left at None, the session decides. A single process loads whole. A distributed ModelOptSession decides with the new resolve_distribution_plan: the model's spec first, which for several models is what supplies the expert-parallel plan, then decide_distribution_plan ([2/20]) with the MODELOPT_EP_DEGREE / MODELOPT_FSDP_SIZE knobs (modelopt.torch.utils.env, moved here from [20/20]). It refuses a plan that would leave an expert-parallel MoE whole on every rank. A caller that loads a whole model its own way decides first: resolve_distribution_plan returns the plan and the config the load must use.
• 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).
• 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.
  • The FSDP2 load refuses a quantized checkpoint, whose tensors it would read as stored, and a partial fsdp_size without expert parallelism, which it cannot shard.

Usage

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

with modelopt_session():
    model = load_pretrained("Qwen/Qwen3-30B-A3B")  # experts sharded over the 8 ranks
    # ... mtq.quantize(model, ...), export_hf_checkpoint(model, ...)

Testing

• New tests/unit/torch/utils/test_load_pretrained_distribution.py (36 cases, CPU, no process group; stand-in loaders record their calls):
  • the knobs, read at call time;
  • the decision: none outside a distributed session, an EP-capable MoE planned expert-parallel, the spec applied first, the refusal to replicate an expert-parallel MoE, knob and argument overrides;
  • each load: the arguments it gets, the config prepared once (a model's preparation can patch its modeling code), and every refusal raised before any loader runs.
  • Mutation-checked: dropping the rank-GPU default, the MTP refusal, the prepared-config handoff, the partial-FSDP refusal or the quantized refusal each fails a test.
• New tests/gpu/torch/utils/test_load_pretrained_distribution.py, on every GPU of the node, load_pretrained left to decide: a dense model loads whole on each rank; with MODELOPT_FSDP_SIZE=world, FSDP2-sharded; a Qwen3-MoE with its experts sharded. Each rank's weights match the checkpoint, and the load records its unplaced keys. Not run before opening this PR; CI's GPU workflow runs it.
• All of tests/unit on this branch, in CI's linux-lane environment (Python 3.12, torch 2.14, transformers 5.18, CPU): 5274 passed, 27 skipped, 12 xfailed, none failed. The Sphinx docs build is clean.
• pre-commit, mypy included: clean on the branch's diff against main.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ load_pretrained keeps its behavior outside a distributed session; distribution and resolve_distribution_plan are new.
• 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.48.0, Misc.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 19 of 20 of expert-parallel PTQ for Hugging Face MoE models; merges after [18/20].

[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]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

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

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

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-13 19:38 UTC

[shengliangxu]

Superseded by #2779 ([3/21], this routing without expert parallelism, now right after #2737) and #2780 ([13/21], the expert-parallel half, after the EP loader and model specs it needs). Moving this branch down the stack made it part of this PR's old base, shengliangx/hfep-13-ep-export, so GitHub marked this PR merged into that branch. Nothing was merged into main.