[shengliangxu]

Note

Re-cut again on 2026-10-29, so more of the series merges in parallel: #2780 and #2747 now build on this PR directly instead of on #2743, and #2743 stays on #2742. This PR's base branch was rebuilt from the current #2779 (now on main), #3067 and #2740. The series first became a tree on 2026-10-27, from main 1521053 (this PR was [9/21]); it was rebased onto main 7c3c868 later that day, after #2739 merged, and onto main f03651e that night.

Base: shengliangx/hfep-base-6a-ep-loader, a branch (not a PR) that merges #2779, #3067 and #2740, the PRs this one builds on; #3067 and #2740 bring #3074 and #3122 below them. The diff below is this PR's alone, and it merges after all of them.

Re-cut: this PR's diff is unchanged, except that _eager_only_experts takes #3067's rule (see the last Testing line), its _auto_class table rows moved to test_model_load_utils.py, where #2737 moved that test, and two docstrings keep main's newer wording. Review fixes since the first re-cut are listed at the end of Testing.

Series layout

• main ([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101 and [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 have merged)
  • [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 calibration on each rank's own experts
    • [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm
      • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
      • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
  • [OMNIML-6054] Route load_pretrained by the distribution plan #2779 load routing
    • [OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741 expert-parallel loader (on [OMNIML-6054] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740) (this PR)
      • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
        • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
      • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
        • [OMNIML-6076] Describe expert parallelism in the model specs #2746 expert parallelism in the model specs
          • [OMNIML-6078] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters #2748 Kimi-K2.5 and Kimi-VL adapters (on [OMNIML-6076] Describe expert parallelism in the model specs #2746 and [OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747)
          • [OMNIML-6079] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter (on [OMNIML-6076] Describe expert parallelism in the model specs #2746 and [OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747)
          • [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] Describe expert parallelism in the model specs #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)
      • [OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747 on-the-fly decompression
  • [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold
  • [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 distributed-writer pieces
    • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
      • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export

What does this PR do?

Type of change: New feature

load_and_prepare_ep(ckpt_path, plan) loads a model by an expert-parallel DistributionPlan (#2736, #2737), with every rank both data-parallel and expert-parallel. It checks the plan against the run before anything else. It records no plan on the model: load_pretrained records the plan once for every load (#2779; the expert-parallel route is #2780), and a caller loading through this loader directly records it with set_distribution_plan, as the GPU tests here do. It lives in modelopt.torch.utils.plugins.model_load_utils with the other loaders. modelopt.torch.utils.moe keeps what expert parallelism is built from: finding the routed experts (#2737), the all-to-all transport (#2738), and the EP mesh calibration reads back.

• transformers 5.7 or later is required, and checked before the load (moe.require_transformers): below 5.6, transformers' expert-parallel load shards nothing, without an error, and 5.6 builds a model only with the experts implementations it ships, so it refuses torch_a2a.
• Experts are sharded as they load. transformers' own expert-parallel load (DistributedConfig(enable_expert_parallel=True)) builds only each rank's experts, so no rank ever holds the whole expert set. The model's base_model_ep_plan is swapped to the torch_a2a styles from [OMNIML-6059] [1/21] Add the DP=EP all-to-all expert transport #2738, so tokens travel to the rank that owns their expert. The swap lasts only for the load: the model keeps the plan it was sharded by as model.tp_plan, which calibration re-installs the forwards from, and the config it keeps goes back to the checkpoint's own plan, so an exported config.json names no ModelOpt-only style. The experts implementation the forwards dispatch on (torch_a2a) stays set; transformers leaves it out of a serialized config.
• The dense part is replicated. Each rank runs its own data through it, so calibration shards the data per rank ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740).
• Calibration's amax sync learns the layout without quantization code knowing how the model was loaded. Each module gets a _parallel_state, which quantization keeps when it converts the module (Megatron's expert MLPs rely on the same rule):
  • dense modules reduce over the whole world;
  • the routed-expert containers reduce over their expert-data-parallel replicas, and share each layer's input amax over the expert-parallel group.
• The routed experts are found from the same base_model_ep_plan that sharded them, not from class names, by the rule the decision ([OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737) counts them with. The loader builds its model through the class the decision chose, except that require_lm_head=False builds a VLM without its head, through AutoModel (the EP degree and the expert containers do not depend on the head). The experts are quantized as they are, through the grouped GEMM of their torch_a2a forward. Experts it cannot quantize that way are refused: those stored transposed (GPT-OSS's), told by their is_transposed flag or, for gated ones that do not say so, by their shapes, which it would leave unquantized; those ModelOpt quantizes through a wrapper of their own (Llama4's), whose forward drops the transport's process group; and those whose fused-experts wrapper quantizes on the eager forward only (FP8's), which the transport does not run. A plan whose expert containers match no module of the loaded model is refused after the load, before any parallel state is seeded. So is a routed-expert container calibration would not take for one by its name (no experts in it, model_calib._is_routed_expert), whose weight amax it would reduce over ranks holding different experts. No quant-module registration is added or removed. The transposed and eager-only rules are the ones quantization itself applies, from one helper each in the quantization HF plugin (_transposed_experts, _eager_only_experts).
• A native MXFP4 checkpoint is dequantized to bf16 as it loads (Mxfp4Config(dequantize=True)), since quantizers attach to bf16 weights.
• A checkpoint in a format the load cannot dequantize is refused once its config is read, before anything loads; transformers' own quantizer would otherwise load it as stored. Supported formats: unquantized; native MXFP4 here; native block FP8 from [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743; compressed-tensors where a model spec decompresses it as it loads ([OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747). Refused: e.g. GPTQ, ModelOpt-format NVFP4, compressed-tensors float-quantized (GLM-4.5/4.6-FP8).
• The checkpoint tensors the model has no parameter for (an MTP head) are recorded as it loads (record_unplaced_source_keys, from [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733), so the export carries them from that record.
• The EP mesh stays recoverable. For a model whose plan is expert-parallel, expert_parallel_mesh(model) returns the mesh the experts were loaded over (transformers keeps it on the model), so calibration can re-install the all-to-all forwards ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740). A model planned expert-parallel that carries no mesh is refused by name (load it with load_and_prepare_ep).
• The experts land on the EP group even on transformers 5.18+.
  • From 5.18, from_pretrained drops its device_mesh argument once distributed_config is given, and builds a world-sized tp mesh of its own.
  • At an EP degree below the world size, that shards the experts over every rank: on 4 GPUs at EP 2, each rank held 2 of 8 experts instead of 4.
  • On those releases the loader sizes tp to the EP degree and hands from_pretrained the EP mesh for the load. Earlier releases honor device_mesh and are untouched.
  • Each DistributedConfig argument is decided by its own field: tp_plan on DistributedConfig where it has the field (5.10, and 5.13 on), else on from_pretrained, never both; tp_size is the EP degree where transformers builds its own mesh.
  • A load that leaves a rank another share of the experts than ep_size gives it is refused, on every rank alike: whole on every rank (a plan whose expert weights no style shards), or sharded over another mesh (a mesh-builder patch that stopped applying). A DTensor (5.16 on) shows its whole count and this rank's share; a plain tensor's share is checked against the count the config names. Weights that show no share at all, as a container owning none itself would, are refused too, rather than passed unchecked.

The loader reads the plan before it resolves the checkpoint or builds a mesh, and refuses one that styles no routed-experts container, or one whose ep_size does not divide the routed experts its config counts, as the decision counts them: the all-to-all transport finds an expert's rank by dividing its id by each rank's share, so every rank must hold as many. A count the config names none of those ways is read after the load off the experts' weights, which transformers keeps whole on the DTensors it shards them into from 5.16 (5.16.0 splits an uneven count without an error); a load whose weights show no count ep_size divides is refused, as the decision refuses a count it cannot read. It loads in the config's dtype, builds the VLM class with the caller's trust_remote_code, and reads a quantized language model's format from llm_config too.

Usage

from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan, size
from modelopt.torch.utils.moe import expert_parallel_mesh
from modelopt.torch.utils.plugins.model_load_utils import load_and_prepare_ep

# under torchrun, inside distributed_run()
plan = DistributionPlan(size(), ep_size=size())  # size() must divide the routed experts
model = load_and_prepare_ep(ckpt_path, plan, hf_config=config)
set_distribution_plan(model, plan)  # load_pretrained records it; a direct caller does
ep_mesh = expert_parallel_mesh(model)  # for distributed_forward_loop(..., ep_mesh=ep_mesh)

Testing

• New tests/unit/torch/utils/test_ep_loader.py, on CPU:

  • a plan the loader cannot carry out refused before anything is read: one that shards no experts, one for another world size, checked on a real one-rank gloo group, and an ep_size that does not divide the experts (3 over 2 ranks); and, after the load, a plan whose expert containers match no module of the model, containers whose names calibration does not take for routed experts (routed), and a count the config does not name, read off the experts' weights: a DTensor count ep_size divides loads on, one it does not, or plain weights that show none, are refused;
  • the class a VLM's backbone builds through when its LM head is not required (test_model_load_utils.py);
  • routed experts with a quant wrapper of their own refused, with the registry left as it is and the shared experts' wrapper no bar; routed experts under ModelOpt's generic fused-experts wrapper, whose forward keeps the expert-parallel transport, accepted, and under its eager-only FP8 wrapper refused;
  • experts stored transposed refused, whether they say so or only their shapes do, and standard-layout fused experts accepted: the grouped quantization would leave transposed ones unquantized;
  • a checkpoint format the load cannot dequantize (GPTQ, ModelOpt, compressed-tensors) refused before anything is read;
  • transformers 5.6 refused before the load. Its from_pretrained checks a model's experts implementation against a fixed list without torch_a2a; checked against the 5.6.0, 5.6.1 and 5.6.2 releases, while 5.7 and later check the registry torch_a2a is registered in;
  • parallel-state seeding, including a sibling whose name extends an expert container's;
  • the config the model keeps carries the checkpoint's own base_model_ep_plan after the load, not the swapped torch_a2a styles, for a plan on the config and for a class default, each failing without its half of the restore;
  • a rank holding another share of the experts than ep_size gives it refused after the load: sharded over another mesh (sharded-over-another-mesh), or plain weights left whole while the config names the count (never-sharded); and experts whose container owns no weights itself, which show no share at all (owning-no-weights);
  • a tiny Qwen3-MoE built through the loader: model.tp_plan keeps the torch_a2a styles once the config's plan is put back, and the config serializes no experts implementation (checked on transformers 5.7, 5.13, 5.16 and 5.18);
  • the checkpoint's quantization method from the config, a VLM's text config, or a transformers config object;
  • on transformers 5.18+, from_pretrained's own mesh builder handing back the EP mesh during the load, and restored afterwards even when the load raises;
  • the DistributedConfig and from_pretrained arguments for each field/builder combination of the releases from 5.6 to 5.19, and for the builder without a tp_plan field, which no release has and which raised TypeError before the fix;
  • expert_parallel_mesh: the model's mesh for an expert-parallel plan, none for a model with no plan or a replicated one, and an expert-parallel model without a mesh refused.

• Mutation-checked: each test fails on at least one deliberate break of the code, the plan checks and the plan-read mesh included; the count read off the weights (dropping it fails the not-dividing and not-shown rows, dropping its none-shown refusal the latter) and the experts-name refusal each fail without their fix; the no-match, eager-only and missing-mesh refusals, and the refusal of an ep_size that does not divide the experts, each fail without their fix, and so does the transposed rule's shape test (reading the flag alone fails the shapes row, an inverted shape test the standard row).

• The refusal table's "with the fix" column was measured on an earlier revision of this branch, with [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 on top and at the stack tip, which agree on every case; the runs below say which revision they are from. The rest of the GPU tests run in CI only.

• New tests/gpu/torch/utils/test_expert_parallel_load.py: loads a tiny Qwen3-MoE expert-parallel, with EP over every GPU and over half of them, and checks each rank's share of the experts and the logits against the whole model, and that the load recorded its unplaced keys (none, for this checkpoint, which save_pretrained writes per expert) and its plan, on every module. On a 4-GPU node, 2 passed on this branch: EP 4, and EP 2 with each expert on two ranks. On 1 GPU, both skip.

• On this revision, with [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 on top and at the stack tip: the expert-parallel GPU tests pass on 2 GPUs (8 passed, 4 skipped) and 4 GPUs (12 passed), EP over half the ranks included. Those runs did not record their transformers version: the GPU harness installs it per run, resolved under transformers>=5.5,<5.19, which gave 5.18.0 wherever a run printed it. (This note said 5.9 before, which was wrong.) At the previous revision, on transformers 5.18 at the stack tip, a tiny randomly initialized Kimi-K3 quantized expert-parallel on 2 GPUs loaded its experts as DTensors, passed the share check, calibrated and exported.

• On transformers 5.18, 4 GPUs, before and after the fix:

  • at EP 2, each rank held 2 of the 8 experts on a world-sized tp mesh before;
  • after, it holds 4 on its 2-rank EP group.

• The refusal on GPUs: [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742's quantized expert-parallel training test (test_quantized_expert_parallel_model_trains_like_the_whole_model), before this PR accepted the generic wrapper (an earlier revision of the stack) and on this revision. Before, the NVFP4 cases were refused (Expert-parallel PTQ cannot quantize Qwen3MoeExperts); the FP8 cases already passed:

  Dense layers, EP span, format	2 GPUs: before, with the fix	4 GPUs: before, with the fix
  replicated, all ranks, FP8	pass, pass	pass, pass
  replicated, all ranks, NVFP4	fail, pass	fail, pass
  replicated, half the ranks, FP8	skipped	pass, pass
  replicated, half the ranks, NVFP4	skipped	fail, pass
  FSDP2-sharded, all ranks, FP8	pass, pass	pass, pass
  FSDP2-sharded, all ranks, NVFP4	fail, pass	fail, pass

• Tests that differ only in their inputs are merged into parametrized tables. Each merge is checked by mutation.

• The config refusals (unsupported formats, no plan, a TP-only plan) in one table; the config's dtype and trust_remote_code reaching the load; an llm_config format row; and the GPU load test reading the expert-parallel mesh back from the model.

• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.

• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.

• Review fixes, 2026-10-27: the loader keeps the expert-parallel mesh on the model itself. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.

• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8).

• Rebased over [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's last review fix (2026-10-27): _eager_only_experts, which this PR factors out of the backend normalization, carries [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's lookup for a class the on-the-fly registration has not registered yet.

• Review suggestion, 2026-10-27, last round: _auto_class's docstring says require_lm_head=False builds an encoder-decoder model headless too, pinned by a new table row; rebased over [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's AutoModelForCausalLM fallback (Phi-4-multimodal).

• Review suggestions, 2026-10-27, last round: a model planned expert-parallel without its device mesh is refused with advice to pass the model load_and_prepare_ep returned, not a submodule of it; the loader's notes print through print_rank_0; a comment says why the loader imports inside the function.

• Review suggestion, 2026-10-28: the share check's comment says it reads the containers' weights together and relies on one expert layout per model, as every transformers MoE has.

• Declined review finding, 2026-10-27: require_transformers keeps refusing a pre-release below its floor (5.7.0.dev0 < 5.7). Every floor the series checks (5.7, 5.8, 5.13) is released, transformers main (5.21.0.dev0) passes them, and a pre-release of a floor can predate the feature it gates.

• Declined review suggestion, 2026-10-28, fifth round: raising when no config takes the torch_a2a tag; every PretrainedConfig defines _experts_implementation (transformers 5.5.4 and 5.18), so the model's own config always takes it.

• Re-cut on 2026-10-29 (see the note at the top), rebased commit by commit onto the rebuilt base:

  • _eager_only_experts takes [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067's rule, the registry lookup alone. [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 dropped the non-gated wrapper, so the only wrapper the on-the-fly registration still adds is the generic one, which is not eager-only; the fallback lookup for a class not registered yet, carried over from [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 (the 2026-10-27 line above), is gone. The eager-only wrappers (FP8's, GPT-OSS's, Nemotron-H's) register when their modules are imported. The loader's refusal and the backend normalization keep sharing the helper.
  • The _auto_class table rows moved to test_model_load_utils.py. The docstring keeps main's rule that an is_encoder_decoder config counts whether or not it names its architecture, and register_torch_a2a's keeps main's wording beside the 5.7 floor.
  • Pre-commit, mypy included, passes at this head over the whole change against main. On CPU at this head (torch 2.14), this PR's tests and those of the code it touches (test_model_load_utils.py, test_moe_a2a.py, test_meta_dcp_load.py, test_fused_experts.py, the Nemotron-H tests) pass on transformers 5.18 and 5.5.4, and so does the whole tests/unit/torch/utils directory on 5.5.4.

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 [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ✅ Approved 2026-10-28 at 2604148c5f1, the same diff this PR kept until the 2026-10-29 re-cut. The re-cut changed only _eager_only_experts's rule (now [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067's), two docstrings, and where its test-table edits live (Testing).

Additional Information

Part 9 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2779, #3067 and #2740.

Summary by CodeRabbit

• New Features

  • Added support for loading mixture-of-experts models with experts distributed across devices, including all-to-all expert routing.
  • Added an option to load vision-language models without a language-model head.
  • Added checks for minimum Transformers versions required by supported features.
  • Added expert-parallel device mesh support for distributed model loading.

• Bug Fixes

  • Improved validation of expert-parallel configurations and distributed process groups, with clearer errors for unsupported setups.

[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: cbb1e61
• Latest commit, not reviewed: 0c5af29

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 an expert-parallel model loader with separate dense and expert meshes. It configures torch_a2a routing, quantization handling, Transformers loading, and distributed state. Unit and GPU tests cover loader behavior, expert sharding, and output comparisons.

Changes: Expert-Parallel Loading

Layer / File(s)	Summary
Loader contracts and model selection modelopt/torch/utils/moe.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py	Adds Transformers version and expert-mesh helpers. _auto_class can select AutoModel when a VLM does not require an LM head.
Expert-parallel loading setup modelopt/torch/utils/plugins/model_load_utils.py	Adds load_and_prepare_ep with plan, process-group, and world-size checks. It builds dense and expert meshes, configures Transformers loading and torch_a2a routing, handles MXFP4 dequantization, records unplaced checkpoint keys, and prepares model and calibration parallel state.
Loader validation and distributed output tests tests/unit/torch/utils/test_ep_loader.py, tests/unit/torch/conftest.py, tests/unit/torch/utils/test_meta_dcp_load.py, tests/gpu/torch/utils/test_expert_parallel_load.py	Adds unit tests for loader validation, expert wrapper handling, mesh selection, quantization lookup, and mesh-builder restoration. GPU tests verify expert shards, checkpoint metadata, and sharded-model outputs. The one-rank Gloo fixture is shared through conftest.py.

Priority: Normal

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant Loader as load_and_prepare_ep
  participant ProcessGroup
  participant DeviceMesh
  participant Transformers
  Caller->>Loader: Provide checkpoint and expert-parallel plan
  Loader->>ProcessGroup: Check initialization and world size
  Loader->>DeviceMesh: Build dense and expert-parallel meshes
  Loader->>Transformers: Load checkpoint with EP mesh and configuration
  Transformers-->>Loader: Return loaded model

Loading

Merge Risk: Low · up to cbb1e

Sequential EP and ordinary Llama4 workflows can produce incompletely quantized models. Isolate the EP wrapper exclusion before merging, or accept this bounded risk with a separate-process workaround.

Pre-merge checks: 5

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 two modelopt Python files and no examples. The diff adds trust_remote_code with default False and passes that caller value to Transformers; it does not add `trust_remote_…
Title check		The title clearly and concisely describes the primary change: loading MoE models with expert-parallel sharding at load time.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The PR changes only two modelopt Python files and no examples. The diff adds trust_remote_code with default False and passes that caller value to Transformers; it does not add trust_remote_code=True. No added torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), external-input eval()/exec(), or # nosec appears. pyproject.toml and requirements.txt are unchanged.

✨ 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.}

[github-actions]

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

[shengliangxu]

/ok to test 5c154cc

[codecov]

Codecov Report

❌ Patch coverage is 86.04651% with 24 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.75%. Comparing base (a713395) to head (ee17253).

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

Additional details and impacted files

@@                              Coverage Diff                               @@
##           shengliangx/hfep-5-distributed-calibration    #2741      +/-   ##
==============================================================================
+ Coverage                                       61.69%   61.75%   +0.05%     
==============================================================================
  Files                                             649      649              
  Lines                                           72627    72795     +168     
==============================================================================
+ Hits                                            44808    44953     +145     
- Misses                                          27819    27842      +23     

Flag	Coverage Δ
unit	61.75% <86.04%> (+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 e1ef1d1

[shengliangxu]

/ok to test 288db20

[shengliangxu]

/ok to test ac84840

[shengliangxu]

/ok to test 0c5af29