[shengliangxu]

Note

The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [14/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.

Base: #2736. The diff below is this PR's alone, and it merges after #2736.

Re-cut: the gloo_world_of_one fixture is added to tests/unit/torch/conftest.py here (#2741 brought it before), in the same text #2751 adds, so either order merges cleanly. Review fixes since are listed at the end of Testing.

Series layout

• main
  • [OMNIML-6073] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm
    • [OMNIML-6073] 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-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load (merged)
  • [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives
    • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
    • [OMNIML-6081] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold (this PR)
    • [OMNIML-6082] Add the pieces of a distributed HF checkpoint writer #2751 distributed-writer pieces
      • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] Add the pieces of a distributed HF checkpoint writer #2751, [OMNIML-6081] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] Route load_pretrained by the distribution plan #2779 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
    • [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 plan decision
      • [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)
          • [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-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747 on-the-fly decompression
                  • [OMNIML-6078] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters #2748 Kimi-K2.5 and Kimi-VL adapters
                  • [OMNIML-6079] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
                  • [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 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)

What does this PR do?

Type of change: New feature

The fused-experts export splits each 3-D expert weight into full per-expert tensors before it quantizes them. Under FSDP2 that first all-gathers the weight, so every rank ends up holding every expert: a 235B MoE OOMed at about 157 GB per rank, however many nodes it ran on.

• Folding in place. With keep_fused_experts set on the export context, the handler folds each expert's quantizer into the fused weight in place:
  • On a sharded weight it quantizes only this rank's experts, finding each one's quantizer by its global expert number when the module keeps every expert's quantizers (FSDP2), or by its position in the shard when it keeps only its own (expert parallelism). Any other count is refused on every rank before anything is folded.
  • Where a rank keeps only its own experts' quantizers (expert parallelism), the fold first gathers each projection's quantizer states over the mesh dim that shards the experts, so every rank folds, leaves or refuses the projection alike, and a refusal names the global experts. The gather runs on that path only: two small object gathers per fused-experts module.
  • It keeps the result sharded with the original expert count. Uneven splits, such as 3 experts over 2 ranks, keep their numbering.
• Unquantized experts. A projection whose experts are all unquantized stays as it is, which is what the per-expert export writes for it, unless its weights are native FP8 (below). One whose experts mix quantized and unquantized cannot be one fused tensor and is refused before anything is folded.
• Splitting afterwards. split_fused_experts_state_dict later cuts the folded weight and its stacked scales into exactly the per-expert keys the single-process export writes. It takes one module's shard at a time, not a consolidated state dict: its input keeps the fused tensors while the per-expert copies are made.
• FSDP2 windows. When the model's distribution plan ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) FSDP2-shards its dense part, the fold and every dense weight's packing on this path run inside an FSDP2 weight window, opened on the innermost unit that encloses the weight's module. That unit is found by walking up the module's name through name maps built once per export (ExportContext.module_names), instead of scanning the model per module. The window hands FSDP2 the packed weights and reshards on the way out. Inside held_fsdp2_weight_windows(ctx), which the quantized-module export enters, a unit's window stays open across its modules and closes when the next unit's opens or the block exits; elsewhere each module's handler opens and closes its own.

  • Without it, the next reshard puts FSDP2's own unquantized shards back.

  • Opening the window on the unit rather than the root matters for expert parallelism, which wraps only the decoder layers.

  • Time to open the window around each of 10 linears per FSDP2 unit, in a synthetic model of about 1,000 modules per unit (CPU, one gloo rank, torch 2.14):

    Units	Scanning the model per module	Walking up the name
    48	22.5 s	4.8 s
    96	79.8 s	9.5 s

• NVFP4 on a sharded weight.
  • The per-block scale avoids index_put_, which DTensor cannot dispatch.
  • _cast_fp4 packs the local shard and keeps the global shape.
  • The nibble pack uses torch.bitwise_left_shift: x << 4 on a DTensor returns x unchanged (torch 2.11), which packed unrelated codes next to correct scales.

Nothing sets keep_fused_experts yet. The writer that consumes the fold arrives in #2751 and #2752.

The fold drops the source checkpoint's expert scales (*_scale_inv, *_activation_scale) as the per-expert export does, and refuses on every rank a layout that leaves a rank no experts, a fused weight sharded on a dim other than the expert axis or over more than one mesh dim, and an unsharded weight whose num_experts, expert rows and weight quantizers disagree. The shard checks run for every sharded projection, quantized or not. It also refuses native FP8 expert weights in every projection, quantized or not, since it drops each projection's scale_inv (dequantize them first, unpack_weight()): a native-FP8 projection whose experts a recipe leaves unquantized would otherwise be kept as raw FP8 bytes without it. The split cuts a packed or block scale in proportion to the weight rows it cuts, so INT4 blockwise, W4A8 AWQ and FP8 2-D blockwise gate and up get their own scales, and refuses a block scale a gate|up cut would split (128-row blocks over gpt-oss's 2880-row gate, say) instead of rounding.

Usage

N/A: internal to the distributed export (#2752).

Testing

• New tests/unit/torch/export/test_fused_experts_keep_fused.py, on CPU under gloo, one and two ranks:

  • fold-then-split against _export_fused_experts, byte for byte, for gated and ungated experts in FP8, FP8 weight-only and NVFP4;
  • the uncalibrated-expert fallback, unquantized projections exported as is, and a mixed projection, native FP8 weights or a num_experts short of the weight's experts refused with the module untouched, a native-FP8 projection whose experts are all unquantized included;
  • uneven and sharded experts;
  • two ranks folding their own experts under global indices, and with expert-parallel quantizer lists (each rank keeps only its own) against the single-process fold; a mismatched count refused on both ranks; rank 1 disabling some or all of its experts' quantizers refused on both ranks, naming global indices, and every expert disabled left unquantized on both;
  • a fused weight sharded off the expert axis ((Shard(2),)) or over two mesh dims ((Shard(0), Shard(0))) refused, on one gloo rank;
  • the split refusing a block across the gate|up cut, and uneven blocks;
  • NVFP4 packing of a sharded weight against the plain one, and the NVFP4 handler packing inside an FSDP2 layer with the root wrapped and unwrapped;
  • the export handler's switch, and the fused experts folded inside a weight window exactly when the plan shards the dense part.

• Mutation-checked: each test fails under at least one deliberate break of the code, the handlers' two plan reads, the shard-local quantizer pairing, the quantizer-count, native-FP8 (an all-unquantized native-FP8 projection's too), num_experts, placement and split refusals, the ranks' joint decision on a projection, and the window's unit included.

• The export comparisons are one table, the one-rank and faked uneven shards fold into the real two-rank job, and the NVFP4 zero-block tests move beside main's scale-clamping tests. Each drop is checked by mutation.

• The per-expert comparison covers source scales, and INT4 blockwise, W4A8 AWQ and FP8 2-D blockwise experts at 128-wide dims; one expert over two ranks is refused on both.

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

• GPU tests were not re-run on this revision: the fold's gather has run on two gloo ranks only, not yet over NCCL in an expert-parallel export on GPUs.

• 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 fold reads each expert through the module's own accessor, in the weight domain the per-expert export packs; one FSDP2 weight window per unit rather than per linear; a fused-experts module not named experts is refused before anything is folded. 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.

• Review fixes, 2026-10-27, round 2: a unit's FSDP2 window is held only inside held_fsdp2_weight_windows, which the quantized-module export enters, and each handler packs its module in one call (the BMM experts, which pack two weights, refuse the keep-fused export); the fold puts its stacked scales on the weight's device; the fold hands the expert accessor the module with only the projection swapped, which [OMNIML-6073] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's accessor needs (with [OMNIML-6073] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged in, 26 keep-fused and fused-experts unit tests fail without it and pass with it). 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.

• Review fixes, 2026-10-27, round 3: the fold writes each packed expert into one preallocated stack, so it holds the packed projection once (table below); a zero NVFP4 block scale is replaced with a scalar rather than a tensor of ones; a window on the model itself reuses the module names the export context holds; comments say what orders an expert-parallel rank's quantizers and which experts a rank folds inside an FSDP2 weight window. The byte-for-byte comparisons against _export_fused_experts pass unchanged; pre-commit, mypy included, and this PR's own unit tests pass at this head. Peak memory the fold adds over its input, one gated module (hidden 1024, intermediate 1024), CPU, torch 2.14:

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

• Review suggestions, 2026-10-27 evening: the keep-fused fold refuses a projection whose weight quantizers use a GGML format, before anything is folded.

  Experts, format	Packed gate_up	Before	After
  64, FP8	128 MiB	274 MiB	166 MiB
  64, NVFP4	64 MiB	195 MiB	149 MiB
  128, FP8	256 MiB	826 MiB	280 MiB
  128, NVFP4	128 MiB	477 MiB	217 MiB

• Review suggestion, 2026-10-27 night: _local_expert_weight's docstring says which reads of the projection the local rows stand in for. The keep-fused tests' ungated case runs transformers' real NemotronHExperts, so it passes whichever wrapper quantizes Nemotron-H's experts ([OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 gives them one of their own).

• Review fixes, 2026-10-27, last round: the BMM experts (GptOssExperts, Llama4TextExperts) refuse the keep-fused export: a GPT-OSS checkpoint without a quantization_config, loaded with FSDP2 and exported without gathering, would have had its fused gate_up_proj split on the wrong axis. The new test fails without the change. Docstring and wording fixes beside it.

• Review suggestions, 2026-10-27, last round: comments name why the GGML refusal agrees across ranks and where the replicated input_scale is synchronized; two unreachable None guards are removed.

• Review fix, 2026-10-27, last round: the gloo_world_of_one fixture destroys its process group in a finally, also when its mesh fails to build.

• Review suggestion, 2026-10-28: test_meta_dcp_load.py's module-local gloo_world_of_one is deleted in favor of the shared conftest fixture, the same deletion [OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741 makes.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Single-process and gather exports are unchanged; the fold runs only when the export context asks for it.
• 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 no-gather export's entry lands with [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 14 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2746.

The Codex re-review item "Validate quantization consistency across EP ranks before folding" is handled in this PR, by the fold's gather, as well as by #2753's model-wide pre-check.

Summary by CodeRabbit

• New Features
  • Distributed model exports can retain fused expert weights and convert them into per-expert checkpoint entries afterward, while preserving sharded layouts.
• Bug Fixes
  • Improved export handling for sharded expert weights and scales, including uneven expert distributions.
  • Fixed NVFP4 quantization and packing for sharded weights.
  • Improved NVFP4 scale handling for zero-valued blocks to keep quantization results consistent.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

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-12a-keep-fused-fold
• File: .coderabbit.yaml
• Commit: d097fd2
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: e98e5c23-468f-400a-9c79-f450b25e16c7

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

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: b728e3a
• Latest commit, not reviewed: d097fd2

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: 46e8cc98-b044-49ee-9cb5-216f53d15c90

📥 Commits

Reviewing files that changed from the base of the PR and between a3e2acd and b728e3a.

📒 Files selected for processing: 1

• tests/unit/torch/export/test_fused_experts_keep_fused.py

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

Walkthrough

The pull request adds an opt-in path for no-gather export of fused MoE experts. It quantizes fused weights, supports DTensor shards, and adds utilities to split fused state dictionaries into per-expert entries.

Changes: Distributed fused MoE export

Layer / File(s)	Summary
Wire fused export into handlers modelopt/torch/export/registry.py, modelopt/torch/export/hf_export_handlers.py, tests/unit/torch/export/test_fused_experts_keep_fused.py	ExportContext adds the keep_fused_experts option. No-gather export uses the keep-fused exporter, with an FSDP2 weight-update window for dense FSDP. Handler tests cover opt-in behavior and update-window ordering.
Quantize fused expert weights modelopt/torch/export/moe_utils.py, modelopt/torch/quantization/qtensor/nvfp4_tensor.py, tests/unit/torch/export/test_fused_experts_keep_fused.py, tests/unit/torch/quantization/test_nvfp4_tensor.py, tests/unit/torch/conftest.py	Fused expert quantization handles plain tensors and DTensor shards. NVFP4 scale handling and packing support DTensor inputs, and zero-amax blocks receive scale 1. Tests cover quantization configurations, uncalibrated experts, and sharded NVFP4 behavior.
Split fused state dictionaries modelopt/torch/export/moe_utils.py, tests/unit/torch/export/test_fused_experts_keep_fused.py	The splitter converts fused gated and ungated projections into per-expert entries. Rank-local splitting uses shard offsets and expert-parallel indexing. Tests cover scale propagation, unrelated keys, and distributed expert indices.

Priority: Normal

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant ExportContext
  participant HFExportHandler
  participant FSDP2WeightUpdateWindow
  participant KeepFusedExporter
  ExportContext->>HFExportHandler: enable keep_fused_experts
  HFExportHandler->>FSDP2WeightUpdateWindow: pack weights within window when dense FSDP is active
  FSDP2WeightUpdateWindow->>KeepFusedExporter: export fused expert weights

Loading

Merge Risk: Moderate · previous review b728e

Keep-fused export can select the wrong expert quantizer on an EP shard, and can fail when a rank owns no experts. Resolve these distributed-export cases before merging unless their limitations are explicitly accepted.

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		No listed security anti-pattern was introduced. The reviewed Python additions contain no torch.load(..., weights_only=False), numpy.load/np.load with allow_pickle=True, hardcoded trust_remote_code=Tru…
Title check		The title clearly and concisely describes the main change: folding fused expert quantizers into weight shards locally. It matches the pull request objectives and affected export logic.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern was introduced. The reviewed Python additions contain no torch.load(..., weights_only=False), numpy.load/np.load with allow_pickle=True, hardcoded trust_remote_code=True, eval/exec calls, or # nosec comments. The pull request does not change pyproject.toml or requirements.txt.

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

[github-actions]

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

[shengliangxu]

/ok to test 52b4592

[codecov]

Codecov Report

❌ Patch coverage is 96.89922% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 62.11%. Comparing base (f337bf6) to head (cf437ff).

Files with missing lines	Patch %	Lines
modelopt/torch/export/moe_utils.py	96.33%	8 Missing ⚠️

Additional details and impacted files

@@                          Coverage Diff                          @@
##           shengliangx/hfep-9-ep-model-specs    #2750      +/-   ##
=====================================================================
+ Coverage                              61.91%   62.11%   +0.19%     
=====================================================================
  Files                                    655      655              
  Lines                                  73164    73416     +252     
=====================================================================
+ Hits                                   45301    45601     +300     
+ Misses                                 27863    27815      -48     

Flag	Coverage Δ
unit	62.11% <96.89%> (+0.19%)	⬆️

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 d314e75

[shengliangxu]

/ok to test dbb570a

[shengliangxu]

/ok to test d097fd2