[shengliangxu]

Note

[2/3] of #2741, split on review: #2741 grew past the 500-line budget, so it is now three PRs, merged in order: #3131 [1/3] (the helpers, on main), this PR (the expert-parallel loader), and #2741 [3/3] (its post-load refusals).

Base: shengliangx/hfep-base-6a-ep-loader, a branch (not a PR) that merges #2779, #3067 and #2740 on main, the PRs this one builds on (#3131 [1/3] has merged). The diff below is this PR's alone, and it merges after all of them.

Series layout

• main ([OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122, [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 [1/3] and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 have merged)
  • [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 → {[OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067, [OMNIML-6071] Calibrate a distributed model on per-rank data #2740}; [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 → [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779
  • shengliangx/hfep-base-6a-ep-loader = [OMNIML-6054] [2/2] 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 + [OMNIML-6071] Calibrate a distributed model on per-rank data #2740
    • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [2/3] expert-parallel loader (this PR)
      • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3] the loader's post-load refusals
      • [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 (merges after [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3])
      • [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 on-the-fly decompression

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, with every rank both data-parallel and expert-parallel.

• Experts are sharded as they load. transformers' own expert-parallel load (DistributedConfig(enable_expert_parallel=True)) builds only each rank's experts, over the ep submesh of an (edp, ep) mesh.
  • The model's base_model_ep_plan is swapped to the torch_a2a styles for the load, so tokens travel to the rank that owns their expert.
  • The config the model keeps gets the checkpoint's own plan back, so an exported config.json names no ModelOpt-only style.
  • The dense part is replicated, and every rank calibrates on its own data ([OMNIML-6071] Calibrate a distributed model on per-rank data #2740).
• The experts land on the EP group on transformers 5.18+ too. Those releases build a world-sized mesh of their own in from_pretrained; the loader hands them the EP mesh for the load. Each DistributedConfig argument goes where the release has a field for it.
• Calibration's amax sync learns the layout from the model. Each module gets a _parallel_state: dense modules reduce over the whole world, and routed-expert containers over their expert-data-parallel replicas. The routed experts are found from the same base_model_ep_plan that sharded them.
• What is checked before anything loads:
  • the plan and the process group;
  • that every rank is bound to its GPU, and that no rank lacks one while others have theirs (review fixes);
  • transformers 5.7 or later (require_transformers, [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131);
  • the checkpoint's format: unquantized. MXFP4 is refused (review fix): transformers dequantizes it into GPT-OSS's transposed expert layout whatever the model's, so a standard-layout MoE would load its experts transposed;
  • that the plan styles routed experts, and that ep_size divides the expert count, read from the config that holds the plan, as the distribution plan's decision reads it (review fix).
• The model keeps its EP mesh for calibration (moe.expert_parallel_mesh), and records the checkpoint tensors it has no parameter for, for the export.
• A rank that fails alone fails every rank (review fix). Each rank checks its own transformers, reads or is handed its own config, sees the checkpoint where rank 0 made it local, and loads its own share, any of which can fail on one rank alone. Each of those steps runs inside a _refused_together block ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779), whose failure-carrying all_reduce raises the failure on every rank before the next collective:
  • the checks;
  • the checkpoint, made local through ensure_local_checkpoint;
  • the load, through the routed experts' lookup. This rests on from_pretrained running no collective, as transformers' expert-parallel load does not on 5.7 through 5.18: before 5.16 it slices each rank's experts into plain tensors, and from 5.16 it builds DTensors locally (run_check=False, src_data_rank=None). The other ranks finish the load and meet a rank that failed in it in the gather.

Why a loader of its own.

• Not transformers' expert parallelism alone. Its stock styles (ep_router, moe_tp_experts) keep every rank on the same tokens. Each rank runs all of them through its own experts, and an all-reduce over the expert-parallel group sums the results. Calibration needs each rank on its own data, so this load keeps transformers' expert sharding (grouped_gemm) and swaps the routing for the torch_a2a styles. Those send each token to the rank that owns its expert by all-to-all, around the grouped GEMM that ModelOpt quantizes.
• Not the meta/FSDP2 loader. parallel_load_and_prepare_fsdp2 shards whole decoder layers, experts included, and gathers each layer's weights on every rank in every forward. It cannot keep a rank's experts resident and route tokens to them. Expert sharding also cannot start from a meta-initialized model: transformers' GroupedGemmParallel.shard_tensor slices materialized tensors. So this load reads real weights through from_pretrained over the EP submesh.
• The mesh-builder override. From transformers 5.18, from_pretrained drops its device_mesh argument once distributed_config is given, and builds a world-sized mesh from the config's sizes. Whenever the EP degree is below the world size, that mesh shards the experts over every rank instead of the rank's EP group. _experts_sharded_over hands transformers' builder the EP submesh this load built, for the duration of the call only. Older releases honor device_mesh and need no override.

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)
ep_mesh = expert_parallel_mesh(model)  # for distributed_forward_loop(..., ep_mesh=ep_mesh)

Design notes

Moved here from the code comments (review).

• Transport. The load swaps the model's base_model_ep_plan to the torch_a2a styles. The router emits global expert ids unmasked. The experts forward sends each token to the rank that owns its expert by an NCCL all-to-all, runs the quantizable grouped GEMM, and sends the result back by a second all-to-all, instead of every rank computing every token and all-reducing. Every rank sees different tokens, hence per-rank calibration data. Nemotron-H's logits-only gate has no ep_router entry, so its nemotron_experts_ep also maps to torch_a2a_experts.

• Meshes. Two factorizations of the same world:

  • (edp, ep) for the experts: sharded over ep, replicated over edp. The ep submesh also drives from_pretrained's expert sharding and the transport.
  • (dp,) for the dense part, replicated over the whole world.

  The all-to-all case is edp == 1, and the dense amax always reduces over the whole world. edp > 1 adds data parallelism across expert-parallel groups, at the cost of per-rank memory, with no concurrent edp and ep collectives.

• Real weights. Expert parallelism cannot be applied to a meta-initialized model, because transformers' GroupedGemmParallel.shard_tensor slices materialized tensors. So the load reads real weights through from_pretrained over the ep submesh.

• Release floor. transformers 5.6 has the MoE expert-parallel framework, but builds a model only with the experts implementations it ships, never torch_a2a. A model whose expert-parallel support landed later states its release in its spec.

• No tensor parallelism. transformers applies expert parallelism through its tensor-parallel path: with enable_expert_parallel=True, model.tp_plan returns the expert plan. tp_plan="auto" is only the switch that makes tp_size > 1 and so reaches apply_tensor_parallelism. The swapped plan keeps expert entries only, so attention and embeddings stay replicated.

  Where tp_plan goes depends on the release: 5.17 and later refuse it beside a distributed_config, and DistributedConfig has no tp_plan field before 5.10, nor in 5.11 and 5.12. Passed alone on 5.17 and later, it rebuilds the config without enable_expert_parallel. So the dataclass fields decide: each argument goes on DistributedConfig where it has a field for it, else to from_pretrained.

• Collective-free load. transformers' expert-parallel from_pretrained runs no collective on 5.7 through 5.18. Before 5.16 it slices each rank's experts into plain tensors, and from 5.16 it builds DTensors locally (run_check=False, src_data_rank=None). So the other ranks finish the load and meet a rank that failed in it in the guard's gather.

• Experts implementation. from_pretrained's experts_implementation argument resolves non-FP8 implementations back to grouped_mm, but the experts forward reads config._experts_implementation when it runs, so the load sets it on every (sub)config. Without that, the torch_a2a_experts style would pass a fifth argument, the process group, to the grouped-mm forward. transformers keeps it as _experts_implementation_internal, which a serialized config leaves out.

• Warmup. The load creates each amax collective's NCCL communicator in the same order on every rank, with a barrier around each, rather than lazily mid-calibration. Lazy creation deadlocks when subgroups reach their first collective in different phases at DP>1. There is no forward warmup, since a variable per-group time would put the ranks out of step first.

• GPU binding. Importing modelopt initializes CUDA wherever torchao and triton are installed: torchao's mx_formats kernels call has_triton(), which reads torch.cuda.current_device(). So in a launch that binds no rank, every rank is already on cuda:0 when the load runs, and DeviceMesh, which binds a rank only while CUDA is uninitialized, leaves it there. On a gloo group, the load refuses such a run on every rank: ranks on the same GPU, compared by UUID. An NCCL launch in that state fails earlier, in init_process_group, where NCCL refuses the duplicate GPU. Where CUDA is still uninitialized, the load refuses ranks whose local rank names no GPU they see.

Testing

• tests/unit/torch/utils/test_ep_loader.py (CPU), covering:

  • the plan, group, floor and config refusals;
  • the Hub ID made local before the weights load;
  • the config's dtype and the class it builds through;
  • each DistributedConfig argument across the 5.6–5.19 field layouts;
  • the 5.18 mesh builder;
  • the restored plan and the saved config;
  • the parallel-state seeding, expert_parallel_mesh and the quantization-method read.

• New: test_a_load_failing_on_one_rank_fails_every_rank runs the load on two real gloo ranks, with the meshes' collective stood in by a bounded barrier. One rank fails alone at each guarded step:

  • an unreadable config;
  • transformers 5.6;
  • a GPTQ hf_config;
  • a checkpoint it cannot see;
  • out of memory in the load;
  • an unreadable weight file.

  Every rank raises that failure. With the guards made rank-local, all six rows fail, the other rank waiting at the barrier (about 8 s each, bounded).

• 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. Unchanged from [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741, where it passed on 2 and 4 GPUs. On this split it ran at 96be72d (below).

• 2026-10-30: pre-commit, mypy included, passes against main. tests/unit/torch/utils and the fused-experts tests pass on CPU at this head, on transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30 (Claude's review): a plan in which some routed-expert container has no grouped_gemm entry under it is refused before anything loads, since grouped_gemm is the only expert-weight style this load shards by (two new table rows fail without the check; every plan transformers 5.18 ships and every ModelOpt spec's plan passes it); the mxfp4 comments no longer name GPT-OSS; the docstring says to pass hf_config on many ranks; __all__ lists is_fsdp2_text_only_vlm. Pre-commit, mypy included, passes against main; the loader, model-loading and MoE unit tests pass on CPU on transformers 5.18 and 5.5.4.

• Rebased 2026-10-30 onto main 99c3834, with [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 rebased there too; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30, round 2 (Claude's review): the load refuses a plan whose routed-expert containers match no module of the loaded model, inside its guard, so calibration cannot reduce the experts' amax over ranks holding different experts; this refusal and its test move here from [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741. It also refuses, before anything loads, a plan that styles a weight under a routed-expert container by a style the load drops; every base_model_ep_plan transformers 5.18 ships and every ModelOpt spec plan passes the check. Both new tests fail without their checks. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18 and 5.5.4.

• 2026-10-30: the no-module plan test's unplaced-keys stand-in returns a list, as the function does (test only). The loader tests pass on CPU at this head, on transformers 5.18 and 5.5.4.

• GPU run, 2026-10-30, at 96be72d (the guards around from_pretrained on NCCL, until then checked on CPU gloo only): tests/gpu/torch/utils/test_expert_parallel_load.py on NVIDIA GB200, torch 2.11.0a0 (nv26.02), transformers 5.9.0, pytest with -p no:rerunfailures.

  GPUs	Result
  2	1 passed, 1 skipped (22 s)
  4	2 passed (27 s)

• Rebased 2026-10-30 onto the rebuilt shengliangx/hfep-base-6a-ep-loader, which after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 merged holds [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 (on [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140, which moves fsdp2_wrap into modelopt/torch/utils/fsdp2_utils.py), [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 on main 62a5c48. This PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU run above was taken at 96be72d; the next entry is one on the rebased stack.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 9c917ae, which contains this PR's head b726fb3 ([OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 adds the dense-sharded and quantized-training cases to the file): tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  2	8	4	0
  4	12	0	0

  On 1 GPU every case skips, so it is no evidence. On 2 GPUs the four cases with expert parallelism over half the GPUs skip; on 4 GPUs all 12 run: each rank's experts and the logits against the whole model, and FP8 and NVFP4 quantized training (scales and gradients), with the experts over every GPU or half of them and the dense part replicated or FSDP2-sharded.

• Review fixes, 2026-10-30 (meenchen's review bot):

  • The expert count is read from the config that holds the plan (4996a25). A new loader test on the decision test's nested config fails without the fix.
  • A load in which some ranks have a GPU of their own and others do not is refused on every rank before either mesh (15e91d9). The bounded two-rank test's new no-gpu row fails without the check.
  • The load's helpers and the test file's transformers imports are module-scope imports (3af30b9).
  • The docstring says what the in-load failure guarantee rests on (e251f6c).
  • Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Review fixes, 2026-10-30, round 2 (meenchen's and Claude's reviews):

  • The load refuses an MXFP4 checkpoint (23feb60). On transformers 5.9.0 and 5.18.0, a tiny Qwen3-MoE loaded from a natural MXFP4 checkpoint got gate_up_proj (4, 64, 192) instead of (4, 192, 64), with no error. The config-refusal table's new mxfp4 row fails without the change.
  • The GPU check is decided from one gather and names the GPU-less ranks (98ee707). The docstring names the transformers releases the in-load guarantee was checked on (1ba50c6).
  • The quantization plugin's TP-linear converter no longer warns about a mesh without TP-sharded Linears on a model recorded expert-parallel (1491b66). Its new test fails without the change.
  • Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. So do the HF quantization plugin's tests.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 7192e05, which contains this PR at e251f6c with the mixed-GPU gather: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  On 4 GPUs every case runs, the four with expert parallelism over half the GPUs included, as at the run before. 1 GPU skips everything, so it is no evidence.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 7c83230, which contains this PR's head 1491b66: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  The same result as the run before, with the same test file, so it speaks for the code underneath. The file loads an unquantized checkpoint on transformers 5.18, so it cannot show the MXFP4 refusal or the block-FP8 floor of 5.8 firing: the config-refusal table's mxfp4 row covers the MXFP4 refusal, on CPU.

• Review fixes, 2026-10-30, round 3 (meenchen's review bot):

  • The GPU check's predicate, which the load routing below defines ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779), now takes the device a rank is bound to once CUDA is initialized, else the one its local rank numbers, as torch's DeviceMesh does. The base is rebuilt on [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 47ca6d1 for it. A new bounded two-rank test (ad9559d) shows each process one GPU, bound to cuda:0 under local ranks 0 and 1: both ranks pass the check and go on to the meshes. It fails with the old predicate. The no-gpu row is unchanged.
  • The tests build their model with the shared model_with_decoder_at from tests/_test_utils/torch/moe_models.py (9a0f5e0).
  • The style swap is defined once, before the checks, and the dropped-style check and the swapped plan both read the styles it keeps (3fc55b4).
  • Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Rebased 2026-10-30 onto main 059588b, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 3b814a6, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 fbff3f3 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 186efae; this PR's own commits are unchanged.

• Review fix, 2026-10-30 (Claude's re-review): a run in which some rank sees CUDA but is bound to no GPU yet is refused on every rank, telling the caller to bind each rank to its GPU first, as distributed_run() does (eb6c382). Before, such a rank was read as having no GPU and the run was refused with "Run one process per GPU"; building the meshes would have failed there anyway, as DeviceMesh binds an unbound rank by its local rank. The one-GPU-per-process test runs that launch bound and not yet bound; the not-yet-bound case fails without the check. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Review fixes, 2026-10-30 (Claude's review at eb6c382):

  • The two-rank GPU-check tests pin, in every row, the CUDA state the check reads, rather than read the host's (69179ec). The gloo workers bind no GPU, so on a host with a visible one a row could be refused before the step it exercises, depending on what earlier tests had initialized. On a host with one visible GPU, a whole-file run at eb6c382 failed the no-gpu row; on a host faked to show GPUs no rank has bound, four rows failed. All pass now.
  • The plan-check comment no longer says packed_* entries are dropped: any other style under a routed-expert container is refused (56353ef).
  • No source behavior changed. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 7540a2e, which contains this PR at eb6c382: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  The same result as the runs before. On 4 GPUs every case runs, the four with expert parallelism over half the GPUs included; 1 GPU skips everything, so it is no evidence. Since that head, [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 changed only tests and one comment, so the result speaks for this head's code.

• Review fixes, 2026-10-30 (Claude's and meenchen's reviews at 56353ef):

  • On a gloo group, the GPU check refuses a run in which two ranks are on the same GPU (57548c6); an NCCL launch in that state fails in init_process_group. Importing modelopt initializes CUDA wherever torchao and triton are installed, so a launch that binds no rank has every rank on cuda:0 at the load, and DeviceMesh leaves it there. The not-yet-bound refusal is narrowed to a rank whose local rank names no GPU it sees.
  • The layout test covers one GPU each (bound, and not yet bound), every GPU not yet bound, and every rank on cuda:0. The narrowing and the new refusal each fail their cases without them.
  • The single-process tests pin no GPU, through a gloo_world_of_one override (b658f37).
  • The docstring keeps the caller contract, and the comments a line or two of why. The rationale is under Design notes (75bc2d6).
  • At 56353ef on a host with one visible GPU, the unit file passed 48 of 48, and the two-rank test alone 7 of 7.
  • Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Rebased 2026-10-30 onto main d464d30, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 8f9d877, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 d452768 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 9a87d15; this PR's own commits are unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU runs, 2026-10-30, at an earlier head with the same own diff ([OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's 6f5391d, which contains this PR at 75bc2d6), on transformers 5.18.0:

  • tests/gpu/torch/utils/test_expert_parallel_load.py:

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  On 4 GPUs every case runs, the four with expert parallelism over half the GPUs included; 1 GPU skips everything, so it is no evidence.

  • tests/unit/torch/utils/test_ep_loader.py on a host with one visible GPU: 59 of 59 passed.

  • The shared-GPU check, on one node with two processes seeing the same two GPUs and neither bound, so that importing modelopt left both on cuda:0:

    Group	Result
    gloo	Both ranks refuse: "rank(s) [0, 1] are bound to the same one".
    nccl	init_process_group itself fails on both ranks ("Failed to initialize NCCL communicator: invalid usage"), before the load.
    nccl, one GPU per process (control)	No refusal; the load succeeds on both ranks.

    So the check catches gloo groups; NCCL launches fail in init_process_group.

• Tests moved 2026-10-31 (meenchen's review): the expert-parallel load's cases from test_ep_loader.py to a section at the end of test_model_load_utils.py, and the two expert_parallel_mesh cases to test_moe_a2a.py, beside moe.py's other tests (7747c5d). The tests themselves are unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils collects and passes the same cases as before on CPU, on transformers 5.18, 5.9.0 and 5.5.4.

• Rebased 2026-10-31 onto main e352cac, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 c78d2e5, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 adb9757 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 dea9b79; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-31, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 939bf6e on main e352cac, which contains this PR's head 930bb0f: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  The same 12 cases as before, now over the rebuilt base, which includes main's fused-MoE GPTQ change. On 4 GPUs every case runs, the dense-sharded ones included: each rank's experts and the logits against the whole model, and both quantized formats' training. 1 GPU skips everything, so it is no evidence.

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] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ❌ Requested. The loader was approved as part of [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 on 2026-10-28; the guards are new.

Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. The loader's review history is on #2741.

[shengliangxu]

/ok to test 84fe2d9

[shengliangxu]

/claude review

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review: 3

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-6a2-ep-loader-core
• File: .coderabbit.yaml
• Commit: 930bb0f
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: c82d5314-8ca3-42d2-8ae8-9e107775747d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• Trigger review

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

[claude]

[IMPORTANT Algorithm] The plan filter recognizes only grouped_gemm as an expert-weight shard style, so a plan that shards the experts another way loads unsharded — silently, with wrong results.

swapped_plan keeps an entry only when v in swap or v == "grouped_gemm", and then tags every routed-experts container torch_a2a_experts, whose shard_tensor is deliberately a no-op (moe.py: return param[...], commented "Expert weights shard via packed_colwise/rowwise + grouped_gemm, so shard_tensor is a no-op here").

So the expert weights are sharded only by the surviving grouped_gemm param entries. Two plan shapes break:

• a plan that styles only the container ({"layers.*.mlp.experts": "moe_tp_experts", "layers.*.mlp.gate": "ep_router"} — exactly the _EP_PLAN the unit tests use, and which the pre-load checks accept): stock moe_tp_experts is what would have sharded, and torch_a2a_experts replaces it with a no-op;
• a plan that shards the expert params with packed_colwise / packed_rowwise beside a container style: the filter drops those entries.

In both cases every rank ends up holding all experts, while torch_a2a_experts_forward routes each token by expert_id // (num_experts // ep_size). Nothing refuses it: the pre-load checks only verify that _routed_expert_keys is non-empty and that ep_size divides the expert count. The result is a calibration that quantizes the wrong experts with the wrong token pools — the worst failure mode, since it produces a checkpoint rather than an error.

Suggested fix: after the load, assert this rank's expert dim is num_experts // ep_size (the docstring already promises "A count the config names none of the ways this reads is checked once the experts load, from their weights" — this is the same check, and it is the one that makes the shard style assumption safe). If that verification is what [3/3] lands, please add the refusal here instead, or at minimum refuse up front a plan whose routed-expert containers have no grouped_gemm param entry under them, so this PR cannot mis-load on its own.

[shengliangxu]

Agreed: this slice must not mis-load on its own. Fixed in 4b6470c with the up-front refusal you suggested.

With the other plan checks, before anything loads, the loader now refuses a plan in which some routed-expert container has no grouped_gemm entry under it. That covers both shapes you named:

• a container-only plan;
• one that shards the expert weights by packed_colwise/packed_rowwise.

Every base_model_ep_plan that transformers 5.18 ships (52 models) and every plan ModelOpt's model specs supply put grouped_gemm under each container, so no real plan is refused.

The test's stand-in plan gains the grouped_gemm entries the real plans have. Two rows of the config-refusal table pin the new check: a container-only plan and a packed_colwise one. Both fail with the check removed.

The post-load check of each rank's actual expert share (#2741 [3/3], on top of this PR) still catches what a plan cannot show, such as experts sharded over another mesh. #2780, the PR that routes users to this loader, merges only after it.

[claude]

[SUGGESTION] This comment and the mxfp4 branch at L1858-1870 contradict each other: here GPT-OSS is "which the load refuses", there the log line reads native mxfp4 checkpoint (e.g. gpt-oss) -> Mxfp4Config(dequantize=True). There is no transposed-layout refusal anywhere at this head (grep transpose in this file finds only this comment), so a GPT-OSS checkpoint is accepted today.

A wrong comment is worse than none here, because it is the only place a reader would look for that constraint. If the refusal lands with [3/3], say so ("refused once the experts load, in a follow-up"); otherwise drop the clause and keep the layout caveat factual.

[shengliangxu]

Right, nothing at this slice refuses transposed experts; that refusal comes with #2741 [3/3]. Fixed in 4b6470c:

• The comment now says only that mxfp4 is for experts stored in the standard layout.
• The mxfp4 branch's comment and log line no longer name GPT-OSS as an example.

[claude]

[SUGGESTION] The load device comes from torch.cuda.current_device(), not from the mesh, so the "device-bound process group" the docstring requires is assumed but never checked. A run that called init_process_group("nccl") without torch.cuda.set_device(local_rank) leaves current_device() == 0 on every local rank, and on 5.18+ releases (the only ones that reach this stub) every rank would place its shard on cuda:0 — an OOM or a silently co-resident load rather than a clear error.

The loader already front-loads the plan / group / transformers / config checks precisely so a bad setup fails before anything loads; this is the same kind of precondition. Consider adding it to the _refused_together("the expert-parallel load's checks") block, e.g. refuse when torch.cuda.current_device() != torch.distributed.get_rank() % torch.cuda.device_count(), or derive the index from the mesh's local coordinate instead of the ambient current device.

[shengliangxu]

Not changed: an unbound device does not lead to the silent co-resident load described here.

• Supported entry point: the loader documents that the caller binds the device. distributed_run, which hf_ptq and the GPU tests go through, binds it with torch.cuda.set_device(local_rank()) before it creates the group.
• Modelopt's group without a binding: the group is cpu:gloo,cuda:nccl, so the guards' object collectives run on the CPU and leave CUDA uninitialized. Torch's init_device_mesh then binds an uninitialized process by LOCAL_RANK, or by rank % devices_per_host, before from_pretrained reads torch.cuda.current_device().
• A caller's NCCL-only group without a binding: every rank's first object collective (guard 1's) lands on cuda:0. NCCL then fails loudly at the load's first collective.

A check of our own would have to guess the launcher's binding rule. Comparing against LOCAL_RANK refuses one-GPU-per-task setups, where every process is cuda:0. Comparing against rank % device_count refuses heterogeneous nodes. Torch's mesh already applies the launcher's own LOCAL_RANK.

[claude]

[SUGGESTION] When the caller passes no hf_config, every rank calls AutoConfig.from_pretrained(ckpt_path) on what may still be a Hub ID — the checks run before ensure_local_checkpoint (L1810), which exists so that only rank 0 talks to the Hub. At a few hundred ranks that is a few hundred concurrent config.json fetches; a 429 or a transient network error on a subset then refuses the whole run (cleanly, thanks to _refused_together, but needlessly).

The ordering itself is well motivated ("a refused load downloads nothing") and the resulting failure is safe, so this is non-blocking. Worth a note in the docstring that callers at scale should hand in hf_config (as hf_ptq does), or read the config once on rank 0 and broadcast it alongside the checkpoint resolution.

[shengliangxu]

Taken as the docstring note you suggested, in 4b6470c. Without hf_config, every rank reads the checkpoint's config itself, from the Hub for a Hub ID, before rank 0 makes the checkpoint local, so callers on many ranks should pass it, as hf_ptq does.

The ordering stays, so that a refused load downloads nothing. Broadcasting rank 0's config would need a config object that loads with weights_only, which PretrainedConfig is not.

[claude]

[SUGGESTION] The new __all__ omits is_fsdp2_text_only_vlm (L92), which examples/hf_ptq/example_utils.py:57 imports from this module — so it is part of the module's effective public surface. Nothing breaks today (nothing star-imports model_load_utils, and from ... import is_fsdp2_text_only_vlm is unaffected by __all__), but per CONTRIBUTING's "__all__ defines the public API" convention the list should match what callers outside the module use, otherwise the next from .model_load_utils import * in plugins/__init__.py silently drops it and autodoc stops listing it.

[shengliangxu]

Done in 4b6470c. __all__ now lists is_fsdp2_text_only_vlm. It was the only name that code outside the module imports from it and the list lacked; the others are load_pretrained, parallel_load_and_prepare_fsdp2, fsdp2_wrap and record_unplaced_source_keys.

[claude]

Claude review — [2/3] expert-parallel loader

Scope: full review (trigger was a bare /claude review). 4 changed files (+1219/-9); reviewed all of them — modelopt/torch/utils/plugins/model_load_utils.py and modelopt/torch/utils/moe.py in full for the changed hunks plus surrounding context, and both test files. Traced the new public symbols (load_and_prepare_ep, expert_parallel_mesh) and the base-branch helpers they lean on (_refused_together, ensure_local_checkpoint, _routed_expert_keys/_routed_expert_modules, ParallelState, distributed_forward_loop, register_torch_a2a, the amax-sync path in model_calib.py).

Findings

CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 4

IMPORTANT — the expert-weight shard style is assumed, not verified (inline, model_load_utils.py:1884-1892). swapped_plan keeps only grouped_gemm as the expert-param style, and torch_a2a_experts.shard_tensor is a deliberate no-op, so the experts are sharded only by the surviving grouped_gemm entries. A plan that styles just the container ({"...experts": "moe_tp_experts", "...gate": "ep_router"} — the shape _EP_PLAN in the new unit tests uses), or that shards the expert params with packed_colwise/packed_rowwise, passes every pre-load check and then loads with every rank holding all experts, while torch_a2a_experts_forward routes by expert_id // (num_experts // ep_size). That calibrates the wrong experts on the wrong token pools and still produces a checkpoint — no error. The docstring already promises a post-load check against the loaded weights; landing it here (or an up-front refusal of a plan with no recognized expert-param shard style) is what makes this slice safe on its own, per the "every sub-PR stands on its own" rule.

SUGGESTIONs: contradictory GPT-OSS/mxfp4 comments with no refusal behind them (1626); the required device binding is assumed but unchecked while the load device comes from the ambient torch.cuda.current_device() (1667); every rank hits the Hub for config.json ahead of ensure_local_checkpoint (1771); __all__ omits is_fsdp2_text_only_vlm, which the examples import (78).

What I checked and found correct

• Collective lockstep. Each guarded step's body is collective-safe: the checks are rank-local before dist.all_reduce; ensure_local_checkpoint's new_group is the block's first statement, and it already broadcasts rank 0's error rather than stranding peers; both meshes are built after the checkpoint block has reduced, so every rank reaches init_device_mesh together. The _sub.size() <= 1 skips are plan-derived and therefore identical on every rank, and the warmup's (dp, ep, edp) order is fixed — no rank-dependent collective anywhere.
• Mesh factorization. DistributionPlan.__post_init__ already enforces world_size % ep_size == 0, so (edp, ep) is always a valid shape, and row-major ordering makes ep_rank == rank % ep_size — which matches both the GPU test's expected shard slice and the transport's expert_id // per_rank owner math.
• Parallel-state seeding. _seed_parallel_states writing module._parallel_state pre-conversion is the right hook: QuantModule.convert only calls _initialize_parallel_state() when parallel_state is None. Dense → world DP, routed experts → edp DP + ep expert group composes correctly with model_calib.py: expert weight amax is excluded from the EP sync by _should_sync_amax_across_ep, and the input-amax share is gated on _share_input_amax_across_ep (set on the HF experts container), so no per-rank shard gets reduced against a peer's.
• Config round-trip. declared_plan is read off vars(ep_cfg) before the swap and restored after (popping to fall back to the config class's default), so an exported config.json names no ModelOpt-only style; _experts_implementation lives in _experts_implementation_internal and stays out of to_dict(). Both are asserted against a real Qwen3MoeConfig.
• The 5.18+ mesh swap. _experts_sharded_over patching transformers.distributed.mixin.initialize_distributed_mesh is verified to actually take effect by driving the real PreTrainedModel.prepare_distribute_model, and the restore is checked on the raising path too — not just asserted in a comment.
• No circular import from moe.py's new top-level from .distributed import distribution_plan (distributed.py references nothing in moe), and distributed_forward_loop really does take the ep_mesh the new docstrings point at.
• Tests. Genuinely good coverage for a loader this version-sensitive: the DistributedConfig field-layout matrix across 5.6-5.19 via synthesized dataclasses, and the six-row two-rank gloo test where one rank fails alone at each guarded step. Both sit in tests/unit legitimately — the multi-rank one uses the session-scoped cpu_dist_workers that CLAUDE.md exempts.

Risk

Moderate. The plumbing is careful and the rank-local-failure work (the new part relative to #2741's approval) is correct and well tested. The single blocking concern is that the load trusts the checkpoint's plan to name grouped_gemm and never confirms the experts actually sharded — on the tested path (Qwen3-MoE) they do, but the failure mode on an untested model family is a wrong checkpoint rather than a crash, which is worth closing inside this slice.

🤖 Generated with Claude Code

[shengliangxu]

/ok to test 4b6470c

[shengliangxu]

/ok to test ab06860

[meenchen]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Nudge: the latest blockers are resolved, but test placement and possible prompt-injection markup in an older review comment need human attention.

Needs action:

• Inspect the embedded </body><parameter name="confirmed">true markup in the supplied Claude comment at model_load_utils.py:1922; confirm it is harmless serialization residue.
• Move loader unit cases into tests/unit/torch/utils/test_model_load_utils.py and MoE-only cases into module-matching tests, following CONTRIBUTING.md’s test-filename rule.

No action needed:

• ✔️ Resolved since the last review: host-independent single-process CUDA fixtures and shortened loader documentation. Earlier bound-device and plan-validation fixes remain present.

[shengliangxu]

/ok to test 7747c5d

[shengliangxu]

Re the review above (test placement and the markup):

• Test placement. Moved in 7747c5d. The expert-parallel load's tests now live in tests/unit/torch/utils/test_model_load_utils.py, as a section of their own at its end. The two expert_parallel_mesh cases are in tests/unit/torch/utils/test_moe_a2a.py, beside moe.py's other tests, and test_ep_loader.py is gone. The tests themselves are unchanged, and the CPU suites collect and pass the same cases as before. The PRs above this one carry their test edits over to the new files.
• The markup. The </body><parameter name="confirmed">true text at the end of Claude's comment on line 1922 is residue from the Claude review bot's own tool call when it posted that comment. It is not an injection, and nothing acted on it.

[shengliangxu]

/ok to test 930bb0f