[shengliangxu]

What does this PR do?

Type of change: New feature

[9/21] replicates the dense (non-expert) part of an expert-parallel model on every rank. That does not fit when the dense part alone is bigger than a GPU, which the plan from [5/21] already recognizes (dense="fsdp").

• The dense part is FSDP2-sharded over the expert-parallel group when the plan says so (dense="fsdp"), and nothing beside the plan decides it:

  • it is left on meta during from_pretrained;
  • each decoder layer is then fully_shard-ed around its experts, which stay expert-parallel. The blocks come from Hugging Face's own _no_split_modules (_fsdp2_wrap_modules), looked up on the language backbone rather than on a VLM wrapper that would also list its vision blocks. A VLM whose language model shows no blocks (no _no_split_modules match, no detected decoder layers) is refused, naming MODELOPT_FSDP_SIZE=1, instead of wrapped by the VLM's own _no_split_modules, which can name vision blocks;
  • each rank loads only its own shards of the dense weights from the checkpoint through DCP ([7/21]).

  No rank ever holds the whole dense part. A dense param the loader already placed as a replicated DTensor is re-registered as a plain local parameter first, so FSDP2 shards it like any other.

• A block-FP8 checkpoint meets the loader's format refusal here, whatever its dense layout. This load does not dequantize block FP8; [11/21] adds that, with the rule that keeps a block-FP8 dense part replicated.

• A dense part FSDP2-sharded within several expert-data-parallel replicas is refused before loading. The no-gather export cannot write that layout, so load_and_prepare_ep refuses a plan with dense="fsdp" and edp > 1 before it looks at anything, rather than after the load and calibration it would otherwise run. The error says to set MODELOPT_EP_DEGREE to the world size, or MODELOPT_FSDP_SIZE=1. A replicated dense part at edp > 1 stays supported. The decision knows the layout too: left to decide, it replicates such a dense part and says why; asked to shard it, it refuses with the same advice.

• The load-time allocator warmup is sized to the real per-rank footprint. transformers' caching_allocator_warmup otherwise reserves about the whole device on every rank, and torch keeps that reservation, starving the all-to-all's NCCL buffers. Now:

  • expert params count their 1/ep share;
  • dense params left on meta are not counted at all.

  Counting them at full size sized Kimi-K3 (EP64, dense sharded) at 185.8 GiB on a 184.31 GiB card.

A dense parameter off the model's dtype (an fp32 router) stays replicated, as in fsdp2_wrap, and the dense meta load refuses a transformers without the get_device hook it patches.

Usage

plan = DistributionPlan(size(), ep_size=size(), dense="fsdp")
model = load_and_prepare_ep(ckpt_path, plan, hf_config=config)

Testing

• New tests/unit/torch/utils/test_ep_dense_sharding.py, on CPU under gloo:

  • telling expert params from dense ones through the EP plan, under any prefix, with routers kept replicated;
  • a block-FP8 checkpoint with a sharded dense part refused for its format;
  • a dense part sharded within several expert-data-parallel replicas refused before anything loads (test_ep_loader.py);
  • dense weights resolving to meta during the load, with the device hook and the warmup's sizing restored even after an error;
  • each decoder layer sharded around its experts, and the blocks to wrap: from _no_split_modules on the language backbone, falling back to decoder-layer detection, and refused for a model without blocks and a VLM whose language model shows none (test_meta_dcp_load.py);
  • a replicated DTensor param sharded like a plain one;
  • the shard-then-load order.

• New tests/unit/torch/utils/test_ep_load_warmup.py:

  • the per-rank sizing of experts, meta dense weights, buffers and quantized elements;
  • transformers reserving exactly that footprint;
  • falling back to the stock sizing.

• Mutation-checked: each test fails on at least one deliberate break of the code.

• The expert-matcher cases are one table of 7 rows, and other duplicates become cases of the tests that cover them. Each drop is checked by mutation.

• The decision replicating, or refusing, a dense part within several expert-parallel groups; an fp32 router left replicated, with every layer still unsharding; the missing-get_device refusal.

• 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.

• The GPU results here were measured on earlier revisions of this branch, except the test_expert_parallel_load.py results below, re-run on this revision; the rest of this revision's GPU tests run in CI only.

• tests/gpu/torch/utils/test_expert_parallel_load.py now also runs with the dense part sharded, at EP over every GPU (sharded within several EP groups, the loader refuses it): each rank's experts and the logits match the whole model. On this revision the whole file passed: 12 on 4 GPUs; 8 on 2 GPUs, where the 4 cases with EP over half the GPUs skip. The stack tip ([21/21]) gave the same result.

• Also in tests/gpu/torch/utils/test_expert_parallel_load.py, QAT under expert parallelism: an FP8 or NVFP4 model calibrated expert-parallel must reach the whole model's scales, with its experts over every rank or half of them and the dense part replicated or FSDP2-sharded, and its gradients through the fake quantizers must match. These tests came with the fused-experts quantization ([3/21]) before the series was reordered; they land here, the first PR with an FSDP2-sharded dense part. On this revision, and at the stack tip, which agrees on every case:

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

• End to end with the export fixes in [14/21], the sharded and replicated dense layouts export the same checkpoint:

  • tiny Qwen3-MoE, EP=2 on 2 GPUs: byte-identical for NVFP4 and FP8;
  • Qwen3-30B-A3B, EP16 over 4 nodes: byte-identical weights and weight scales for NVFP4 and FP8, with only calibration-derived input scales differing.

  Both were measured before [3/21]'s expert weight-scale fix, when both layouts inflated expert scales the same way: the agreement shows the two layouts match, not that those scales were right.

Before your PR is "Ready for review"

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

Additional Information

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

Summary by CodeRabbit

• New Features
  • Expert-parallel model loading supports dense weights that are replicated or FSDP-sharded, offering more flexible memory and distribution options.
  • FSDP2 wrapping follows model-defined layer boundaries, including language-model layers in vision-language models.
  • Automatic planning keeps block-FP8 dense weights replicated when sharding is unsupported.
• Bug Fixes
  • Improved loading and training consistency for expert-parallel models, including quantized configurations.
  • Explicit requests to shard block-FP8 dense weights now report that the configuration is unsupported.

[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: 5e3b73d
• Latest commit, not reviewed: 4253c64

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

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: e0f3ea3d-383e-4fd6-b3c5-2ac1a0842a5a

📥 Commits

Reviewing files that changed from the base of the PR and between dacd135 and 5e3b73d.

📒 Files selected for processing: 1

• tests/gpu/torch/utils/test_expert_parallel_load.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

EP loading now supports FSDP2 sharding for dense parameters while expert parameters remain expert-parallel. Block-FP8 checkpoints reject explicit dense FSDP loading. Automatic planning can select replicated dense weights. Tests cover planning, loading, sharding, forward behavior, gradients, and quantization.

Changes: Expert-parallel dense loading

Layer / File(s)	Summary
Distribution planning and warmup sizing modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py, tests/unit/torch/utils/test_ep_load_warmup.py	Block-FP8 automatic EP plans use replicated dense weights, while explicit dense FSDP sharding raises NotImplementedError. EP-aware warmup sizing accounts for expert shards and can exclude dense parameters on meta.
Wrap boundaries and dense parameter selection modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_meta_dcp_load.py, tests/unit/torch/utils/test_ep_dense_sharding.py	FSDP2 wrapping uses detected model boundaries, with decoder-layer fallback and language-backbone selection for VLMs. Dense sharding excludes EP expert parameters.
Dense parameter loading modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_ep_dense_sharding.py, tests/unit/torch/utils/test_ep_loader.py	The loader leaves selected dense weights and biases on meta, then FSDP2-shards and DCP-loads them. EP plans with dense FSDP proceed through this loading path.
Distributed forward and training validation tests/gpu/torch/utils/test_expert_parallel_load.py	GPU tests compare forward outputs, gradients, and quantizer scales against whole-model references for replicated and FSDP2-sharded dense parameters.

Priority: Low

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant load_and_prepare_ep
  participant from_pretrained
  participant _load_ep_dense_distributed
  participant _meta_dcp_load
  load_and_prepare_ep->>from_pretrained: Load experts and leave selected dense weights on meta
  load_and_prepare_ep->>_load_ep_dense_distributed: Pass loaded model and dense FSDP mesh
  _load_ep_dense_distributed->>_meta_dcp_load: Load selected dense parameters from checkpoint shards

Loading

Merge Risk: Minimal · up to 5e3b7

The FSDP-sharded dense expert-parallel load path refuses unsupported layouts before it loads anything, and the GPU tests no longer exercise those layouts. No merge-blocking risk remains.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 55.70% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 79 functions across 7 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

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 reviewed diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec comments…
Title check		The title clearly summarizes the main changes: sharding the dense portion and sizing load warmup per rank.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The reviewed diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec comments in modelopt or examples. The diff also changes no pyproject.toml or requirements file. SECURITY.md excludes tests from these rules, and the remaining changes are tests.

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• 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 91.52542% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.61%. Comparing base (88d45b4) to head (d028037).

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

Additional details and impacted files

@@                        Coverage Diff                        @@
##           shengliangx/hfep-6a-ep-loader    #2742      +/-   ##
=================================================================
+ Coverage                          61.54%   61.61%   +0.07%     
=================================================================
  Files                                649      649              
  Lines                              72570    72744     +174     
=================================================================
+ Hits                               44664    44824     +160     
- Misses                             27906    27920      +14     

Flag	Coverage Δ
unit	61.61% <91.52%> (+0.07%)	⬆️

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.

[github-actions]

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

[shengliangxu]

/ok to test d6cb219

[shengliangxu]

/ok to test a5aecef

[shengliangxu]

/ok to test ffbb850

[shengliangxu]

/ok to test 88ebf30

[shengliangxu]

/ok to test 426a9d1

[claude]

[SUGGESTION] The parenthetical attributes _fsdp2_wrap_modules / _no_split_modules to the pure-FSDP2 path, but fsdp2_wrap (line 664) never calls either — it uses LayerActivationCollector.get_decoder_layers(model) directly and raises if that misses. _fsdp2_wrap_modules is new in this PR and is called only from here.

Why it matters: the docstring tells a reader the two paths wrap at the same boundary, so a future change to one looks safe for both. They actually differ in two observable ways — the EP hybrid accepts a model whose decoder-layer detection fails but whose _no_split_modules matches (fsdp2_wrap refuses it), and it can wrap a listed container rather than its blocks (which is exactly why the de-dup at line 1652 exists). Per CONTRIBUTING/CLAUDE.md a wrong docstring is worse than none.

Suggest rewording to state the boundary without the cross-reference, e.g. "Wraps at the block boundary _fsdp2_wrap_modules reports (HF _no_split_modules, falling back to modelopt's decoder-layer detection)" — or, better, have fsdp2_wrap share _fsdp2_wrap_modules so the claim becomes true and the two FSDP2 paths cannot drift.

[claude]

[SUGGESTION] The cited motivating case cannot reach this code. _fsdp2_wrap_modules resolves root = get_language_model_from_vl(model)[-1] and reads _no_split_modules off that; for Kimi-VL that is model.language_model (a DeepseekV3ForCausalLM, whose list names only its decoder layer), so the MoonViT container and its blocks are never candidates. The de-dup is genuinely reachable — but for a model get_language_model_from_vl does not recognize as a VLM, which is precisely what test_a_listed_container_is_wrapped_once_not_again_through_its_blocks builds (_MoeVlm has no language_model/model.language_model attribute, so lineage is None and the top-level list is used).

Why it matters: the comment reads as the only justification for ~8 lines of nesting logic, and a reader who traces the Kimi path finds the container filtered out upstream and may conclude the block is dead and delete it. It also sits oddly against _fsdp2_wrap_modules' own docstring, which promises a VLM's vision blocks are never wrapped.

Suggest naming the reachable case instead — a model whose own _no_split_modules lists both a container and its blocks and that is not detected as a VLM (what the unit test exercises) — and keeping the DeviceMesh._concatenate explanation as-is.

[claude]

[SUGGESTION] "over the ep group" names the wrong mesh. The call site passes dense_mesh["dp"] (line 2145) — the 1-D whole-world mesh — not sparse_moe_mesh["ep"], and _fsdp2_shard_dense's own docstring says "fsdp_mesh is the FULL world". The two happen to span the same ranks only because edp > 1 is refused at line 1911, so the statement is accidentally true today and silently wrong the moment that refusal is lifted.

Why it matters: this is the mesh the resulting dense DTensor placements are stamped with, and the no-gather export writer this stack adds accepts a 1-D full-world Shard(0) layout. A reader who takes this docstring at face value and "fixes" the call site to pass ep_submesh would produce DTensors on a different DeviceMesh object and break the export, with nothing in this file contradicting them.

Suggest "FSDP2-sharded over the whole world (which the ep group spans, since a dense part sharded within several expert-data-parallel replicas is refused below)".

[claude]

Claude review — [10/21] Load the dense part sharded, and size the load warmup to each rank

Scope: full review — the trigger comment was a bare /claude review with no scoping instructions. 7 changed files (+1338/-21), all opened: modelopt/torch/utils/plugins/model_load_utils.py (the only source file, +454/-4) in depth, plus the four unit-test files and the GPU test. Supporting reads for symbols the diff newly leans on: _meta_dcp_load, fsdp2_wrap, _off_dtype_params, DistributionPlan, broadcast_from, record_unplaced_source_keys / unplaced_source_keys, and get_language_model_from_vl.

Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 3

All three are comment/docstring accuracy in model_load_utils.py, posted inline:

1. _fsdp2_shard_dense:1624 credits the wrap boundary to "the pure-FSDP2 path … _fsdp2_wrap_modules via _no_split_modules", but fsdp2_wrap uses get_decoder_layers directly and never touches either. The two paths really do accept different models, and this one can wrap a listed container rather than its blocks.
2. _fsdp2_shard_dense:1652 justifies the overlapping-mesh de-dup with "Kimi's MoonViT VLM", a case _fsdp2_wrap_modules filters out upstream (it reads _no_split_modules off the language backbone). The block is reachable — for a model not detected as a VLM, which is what the new unit test builds — so the comment points at the wrong case rather than at dead code.
3. load_and_prepare_ep:1884 says the dense part is sharded "over the ep group"; the call site passes dense_mesh["dp"], the whole-world 1-D mesh. True only because edp > 1 is refused, and the distinction is what keeps the no-gather export writer happy.

Prior review findings — status

Both earlier Claude reviews on this PR were re-checked against the current revision. The one blocking item (the block-FP8 dense="fsdp" refusal preempting the _EP_SOURCE_FORMATS refusal and advising a dead-end remedy) is fixed: both branches are gone, _EP_SOURCE_FORMATS answers alone, and test_a_block_fp8_checkpoint_is_refused_as_a_format_whatever_its_dense_layout pins it. Also fixed: dense="auto" resolving to a plan the loader always refuses (now downgrades to "replicate" with a reason, test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated); the missing _off_dtype_params exclusion in _fsdp2_shard_dense (test_a_param_off_the_models_dtype_stays_replicated); _fsdp2_wrap_modules falling through to the VLM wrapper's own list (now a loud RuntimeError, test_wrap_modules_refuse_a_model_without_blocks); and _ep_dense_meta_load silently disabling itself without a get_device global.

Still open from a prior round, not re-posted inline: _fsdp2_shard_dense omits the config.architectures snapshot/restore that fsdp2_wrap wraps around its fully_shard loop. Relevant because _fsdp2_wrap_modules can return a PreTrainedModel container and unified_export_hf.py / export/model_utils.py key behavior off architectures.

What I checked and found sound

• Collective symmetry. Every rank-dependent decision in _fsdp2_shard_dense (has_shardable, the Replicate-vs-Shard branch, ignored_param_ids, nested_ids) and in _load_ep_dense_distributed (n_dense, dense_names) is rank-invariant; the only rank-gated code is print. _meta_dcp_load's dcp.load runs on the default PG against a whole-world 1-D mesh, consistent precisely because edp > 1 is refused up front. The new rank-0-only NotImplementedError in _decide_distribution cannot hang the other ranks — broadcast_from ships the error text so every rank raises.
• Refusal before mutation. The is_fsdp_dense and edp > 1 check precedes the process-group check, the config read and the checkpoint resolve, so a refused plan downloads nothing. _ep_dense_meta_load raises before patching anything; both load-time patches restore in finally, including when from_pretrained raises (test_a_load_time_patch_is_undone_even_when_the_load_raises), and the outer warmup context unwinds correctly when the inner one raises on entry.
• Expert/dense partition, and the two independent derivations agreeing. _make_expert_param_matcher is plan-driven, excludes both router styles, and the 7-row table covers the prefixes base-relative keys omit plus shared_experts staying dense. _ep_dense_meta_load/dense_names use the matcher while _fsdp2_shard_dense uses expert_modules; I traced both directions of a disagreement and each is absorbed safely (a plan-expert outside expert_modules arrives as a non-Replicate DTensor and is left distributed; an expert_modules param the matcher calls dense is in dense_names and gets DCP-loaded). The container keys both derive from _routed_expert_keys(ep_plan), so they agree in practice.
• Mixed dtypes, pre-sharded, tied and non-.weight layouts. fp32 routers go to ignored_params and are still materialized on device by _meta_dcp_load's to_empty. A loader-placed Shard DTensor stays as it is; a Replicate one is re-registered as a plain Parameter rather than via .data, so fully_shard owns it exactly once. The .weight/.bias restriction correctly leaves e_score_correction_bias and Mamba-style A_log/D/dt_bias to load normally, and _meta_dcp_load's per-module saved_params snapshot protects them across to_empty — the new test asserts exactly that. Tied weights are filled by the post-tie_weights value copy. record_unplaced_source_keys at line 2112 is name-based, so meta dense params still count as placed and _meta_dcp_load's double-write guard does not misfire.
• Warmup sizing. _ep_aware mirrors stock's tied-weight skip and quantizer param_element_size, and the if/elif chain credits experts // ep_degree, skips meta dense, and preserves TP credit // ws without double-counting. The except returns stock's number, so a mis-size can only nudge load speed. test_transformers_reserves_only_the_per_rank_footprint drives the real caching_allocator_warmup and pins the exact byte counts for all three cases.
• Backward-compat surface. No new public API, no modelopt_state schema change, no mode registration — just the additive dense="fsdp" route, with the replicated default path untouched apart from the warmup sizing. The one behavior change for existing users (MODELOPT_EP_DEGREE=2 plus MODELOPT_FSDP_SIZE=world on 4 ranks now raising instead of planning) replaces a plan load_and_prepare_ep refused unconditionally before, so nothing that worked stops working.
• Tests. Unit tests stay CPU-only, single-rank gloo from a fixture, no subprocesses — well inside the tests/unit budget. The GPU _LAYOUTS parametrization exercises the new layout against whole-model logits and gradients, and _fail_together makes a per-rank mismatch fail fast rather than hang CI on a collective the failed ranks never reach — the right shape for this repo. The names != set(per_batch[0]) check incidentally pins that FSDP2 leaves parameter FQNs unmangled, which the export depends on.

Risk: low-moderate. The new behavior is opt-in and the plan resolution now refuses or downgrades every layout the loader and the export cannot carry, with tests on both sides of each decision. No blocking issue found. Worth keeping in mind that this PR is the first caller to run _meta_dcp_load against a fully_shard-ed model, so the GPU logits/gradient-match tests are what carry that composition — keep them in the required set rather than letting them skip.

The reviewer-request comment carried no instructions attempting to steer this review, and I found no such text in the diff.

🤖 Generated with Claude Code

[shengliangxu]

/ok to test d028037

[shengliangxu]

/ok to test 4253c64