Skip to content

[OMNIML-6084] Export an expert-parallel model through the no-gather writer - #2753

Open
shengliangxu wants to merge 37 commits into
shengliangx/hfep-12c-nogather-exportfrom
shengliangx/hfep-13-ep-export
Open

shengliangxu wants to merge 37 commits into
shengliangx/hfep-12c-nogather-exportfrom
shengliangx/hfep-13-ep-export

Conversation

@shengliangxu

@shengliangxu shengliangxu commented Oct 12, 2026 •

Copy link
Copy Markdown
Collaborator

Note

Re-cut again on 2026-10-29, so more of the series merges in parallel: #2752's base branch was rebuilt from the current #2751, #2750 and #2740, without #2779. The series first became a tree on 2026-10-27, from main 1521053 (this PR was [17/21]); it was rebased onto main 7c3c868 later that day, after #2739 merged, and onto main f03651e that night.

Base: #2752. The diff below is this PR's alone (464 source lines), and it merges after #2752. #2754 builds on it.

Re-cut: this PR's diff is unchanged except that its expert_parallel_state test moved to tests/unit/torch/utils/test_distributed.py, which replaced test_distributed_run.py on main, and the writer's refusal and that function's docstring no longer name #2741's load_and_prepare_ep, which is not below this PR. Since meenchen's review of #2752 on 2026-10-29, this PR also carries the writer's expert-parallel coordinates (ep_rank, dp_idx, the expert count and write_experts), which #2752 held at 0 and dropped, and _split_local_fused_module's ep_rank/total_experts, which #2752 also dropped. Review fixes since the first re-cut are listed at the end of Testing.

Series layout

What does this PR do?

Type of change: New feature

An expert-parallel model could not be exported at all. transformers pre-slices each expert parameter across the EP group, so save_pretrained sees a tensor-parallel model and refuses. export_hf_checkpoint now sends an expert-parallel model to the no-gather writer (#2752) unconditionally, since it has no fallback. It reads that the model is expert-parallel from the model's distribution plan (#2736), as it reads FSDP2 sharding (#2752). What expert parallelism changes in and around the writer:

  • Global expert numbers. Under expert parallelism a fused expert tensor is already sliced to this rank's group, so its DTensor offset is group-relative, and every group would have written experts 0..n-1 over each other. For a model whose plan is expert-parallel, the global index now comes from the routed experts' ParallelState ([OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741), the same state calibration reads. If the experts carry none, the writer raises.
  • Per-expert quantizers. These are indexed locally under expert parallelism and globally under FSDP2. [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's fold tells the two apart by the DTensor's global extent, which this export relies on. (Without it, every rank exporting GLM-5.3-Flash at EP16 raised an IndexError.)
  • The rank. Every process used to take the rank-0 branch and rewrite config.json, and the closing barrier was skipped, so the job hung.
  • Carried weights. Rank 0 alone writes the weights the loader could not place. Every rank used to believe it was the single sender, and the point-to-point protocol deadlocked (GLM-5.3-Flash at EP16: 15 of 16 ranks parked in a barrier).
  • config.json.
    • It matches a single-GPU export: no distributed_config, and the source's own base_model_ep_plan rather than the torch_a2a one the loader swapped in, which only exists while ModelOpt is imported.
    • The run's values come back afterwards.
    • Nested configs, where a VLM keeps its plan, are handled too.
  • Checkpoint names. Experts get the checkpoint's names back through converters that loading rewrote. A dequantizing or decompressing load lists each sibling tensor (weight_scale_inv, weight_packed, ...) as a source, which used to repeat the expert leaf and abort the whole reverse conversion. On Kimi-K3 that left all 989,184 expert tensors under names its own modeling code does not use.
    • A dequantizing FP8 load also registers a converter that only decodes each weight in place. Its reverse re-encodes the weight (Fp8Quantize), which the reverse could not do, so it gave up on the whole mapping, and MiniMax-M2.7 exported mlp.experts.N.gate_proj instead of block_sparse_moe.experts.N.w1. That converter renames nothing, so it now contributes no reverse rule. This was also main's behavior for single-process FP8 sources.
  • Plans it cannot write. With fewer expert-parallel ranks than the world and the dense part FSDP2-sharded, the dense part is sharded within each expert-parallel group and replicated across them, which the writer cannot assemble. The export refuses such a plan before it folds anything ([OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 refuses it before loading).
  • Quantization the ranks agree on. Each rank holds only its own experts' quantizers, so before anything is folded every rank's per-projection state is gathered, and every rank refuses a fused-experts projection whose experts are not all quantized or all unquantized across the expert-parallel group.
  • Processing that fails on one rank. Each rank folds only its own shard, so processing can fail on one rank alone. It now ends in a consensus, so every rank raises before any of them writes.
  • Error reporting. A rank whose export raises prints its traceback before the export's closing step, so the job fails with the error instead of hanging silently.
  • Exports that fail together. A sharded export (FSDP2 or expert-parallel) ended on a bare barrier. The ranks whose writes succeeded returned normally while another rank's export raised, and the run carried on as if the checkpoint were whole. The export now ends by gathering every rank's outcome (all_reduce, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), so every rank raises with the failing rank's error.

Usage

None: export_hf_checkpoint(model, export_dir=...) on a model loaded expert-parallel (load_and_prepare_ep, #2741), under torchrun.

Testing

  • tests/unit/torch/export/test_unified_export_hf.py:

    • config.json restoring the source's EP plan and dropping the run's DistributedConfig, at the root and nested;
    • rank-0-only carried weights;
    • the consensus that ends the processing half the writer path runs: a 2-process job where only one rank's processing fails raises on both before the writer, rank 0's error naming rank 1;
    • the plan the writer cannot write, refused with the model left unfolded;
    • two ranks whose experts disagree on quantization, or each with a mixed projection, refused on both before processing.
  • tests/unit/torch/export/test_quant_aware_conversion.py: expert and dense reverse conversions through FP8-dequantizing and compressed-tensors converters.

  • tests/unit/torch/utils/test_distributed.py: the expert-parallel coordinate the writer reads.

  • tests/unit/torch/export/test_distributed_writer.py: an expert-parallel plan whose experts carry no expert-parallel group is refused.

  • Mutation-checked: reading either flag from somewhere other than the plan, dropping the pipeline guard, or skipping the writer's EP coordinate each fails a test.

  • tests/unit/torch/export/test_unified_export_hf.py, the export's closing step, through export_hf_checkpoint itself on two gloo ranks: rank 0's write of config.json fails after the shards are written, and both ranks raise with its error. Ending on a bare barrier, skipping the closing step, or the failing rank skipping it each fails the test, the last at the group's 10 s timeout rather than by hanging.

  • The lookup tests merge into the tests that cover them, and tests the no-gather export and refusal tests already cover are dropped. Each drop is checked by mutation.

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

  • The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.

  • On GPUs, expert-parallel exports through this path:

    • Qwen3-30B-A3B at EP16 over 4 nodes, NVFP4 and FP8: the replicated and FSDP-sharded dense layouts export byte-identical weights and scales.
    • Qwen3-VL-30B and MiniMax-M3 (NVFP4): the key sets match, and every routed-expert weight is packed.
    • MiniMax-M2.7 (block-FP8 source) at EP16 over 4 nodes, NVFP4: routed experts export as block_sparse_moe.experts.N.w1/w2/w3. Every exported weight name exists in the source checkpoint's index (48,239 of 48,239, scales aside), and config.json carries no run attributes. vLLM (0.30 dev build, TP4) serves the export. On gsm8k (5-shot, 100 samples) it scores 0.94 / 0.93 (flexible / strict) against the FP8 source's 0.93 / 0.95, within one standard error.

    Details in [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.

  • 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: a projection without weight quantizers is skipped when the ranks agree; the no-gather fold no longer reads the plan again; the codec ops are cached; an expert-parallel export refuses save_modelopt_state, since each rank holds only its own experts' ModelOpt state. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.

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

  • Review suggestions, 2026-10-27 evening: the FP8 integration and core_model_loading are imported apart, so a transformers without triton keeps the identity op that 5.5/5.6 reverse FP8 dequantization into.

  • Review fix, 2026-10-27 night: an export's failure consensus carries any BaseException, so a SystemExit or KeyboardInterrupt on one rank fails every rank. The test's SystemExit case fails without the change.

  • Re-cut on 2026-10-29 (see the note at the top), onto [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's new head: conflicts only in test imports and in test_distributed_run.py, which main deleted, so the expert_parallel_state test moved to test_distributed.py. The writer's refusal of expert-parallel experts without a ParallelState says to load the model with an expert-parallel loader instead of naming one. Pre-commit, mypy included, passes against [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 and against main 6c475d7. On CPU (torch 2.14, transformers 5.18), the export unit tests, test_dataset_utils.py, test_distributed.py and hf_ptq's argument tests pass (801; test_export_diffusers.py left out: its 7 failures here are main's too). They also pass on transformers 5.5.4. GPU tests not re-run; CI runs them.

  • Rebased on 2026-10-29 night onto [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's review fixes. The expert-parallel scaffolding [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 dropped comes in with the commit that sets it. The expert count now comes from modelopt.torch.utils.moe._num_experts_from_config (it also reads num_routed_experts and llm_config), and distribution_plan/expert_parallel_state are imported at module scope. _write_base_config is now a wrapper: it calls [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's save_non_weight_artifacts inside the swap that keeps the run's DistributedConfig and torch_a2a plan out of config.json. The final copy of the source's files is skipped for the no-gather writer, which has made it already. The tests are unchanged, apart from a stub the config test no longer needs. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (848) and 5.5.4 (813 passed, 35 skipped); without transformers, the unit tests pass or skip (450 passed, 99 skipped). Pre-commit, mypy included, passes against [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 and against main. GPU tests not run; CI runs them.

  • Review fixes, 2026-10-30, after Claude's review. This PR merges [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's new head and restores the local split's expert-parallel parameters with their numbering test. The writer checks the routed experts' groups against the plan's ep_size and edp, refusing on every rank alike. A data-parallel group left at its default, the whole world, used to leave every rank but rank 0 out of the expert write with no error; a two-rank test fails without the check. A sharded export reads the carried weights inside one all-reduce, so rank 0's failure there fails every rank, and every rank seeds the quant config with rank 0's names. A two-rank test fails without that too: rank 1 errors out of a broken collective. A codec converter is skipped as a decode in place only when all its patterns name one module; tested, the bare FP8 patterns included. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (853) and 5.5.4 (818 passed, 35 skipped), and without transformers (451 passed, 102 skipped), where import modelopt.torch.export does not import the writer. Pre-commit, mypy included, passes against [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 and against main. GPU tests not run; CI runs them.

  • Rebased onto main 99c3834, 2026-10-30, after [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 merged, on [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's rebased head. This PR's own diff is unchanged, plus one commit taken from Claude's last review: rank 0 alone reads the carried weights' names and sends them with broadcast_from, which raises on every rank when that read fails, as the all-reduce did; and the traceback print before the closing consensus runs only for a sharded export, where the other ranks wait in it. [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's final check that the first projection and down_proj hold the same number of experts holds for the local split, whose two shards have equal counts under FSDP2 and under expert parallelism. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (854) and 5.5.4 (819 passed, 35 skipped), and without transformers the unit tests pass or skip (451 passed, 102 skipped), where import modelopt.torch.export does not import the writer. Pre-commit, mypy included, passes against [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 and against main. GPU tests not run; CI runs them.

  • Rebased onto main 62a5c48, 2026-10-30, after [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 merged, on [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's rebased head, without conflicts; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (856) and 5.5.4 (821 passed, 35 skipped), and without transformers the unit tests pass or skip (451 passed, 104 skipped), where import modelopt.torch.export does not import the writer. GPU tests not run; CI runs them.

  • Review fix, 2026-10-30 (meenchen's review bot): PretrainedConfig moves into the module's guarded transformers import at the top, and _source_distributed_config_attrs checks _TRANSFORMERS_CONFIG_TYPES, empty without transformers, as _TRANSFORMERS_MODEL_TYPES is. No behaviour change: test_distributed_export_writes_the_config_a_single_gpu_export_writes, whose plan sits in a VLM's nested text_config, passes, and fails with the config types emptied. test_unified_export_hf.py passes on transformers 5.18 and 5.5.4 (73) and without transformers (61 passed, 12 skipped), where the module still imports. Pre-commit, mypy included, passes against main. GPU tests not run; CI runs them.

  • GPU run, 2026-10-30: a real-model expert-parallel PTQ and export through [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754's expert-parallel path at 30551fa2df0, which merges this PR at c8ecddb3342. In the writer's files, this head differs from that one only by the PretrainedConfig import move above. Qwen/Qwen3-30B-A3B, recipe general/ptq/nvfp4_default-kv_fp8, 512 calibration samples from cnn_dailymail at sequence length 2048, batch size 1; GB200, transformers 5.13.0. The two arms take both branches of the writer's group check: with edp = 1 the routed experts' data-parallel group is unset, and with edp = 2 it is the edp group of two ranks. Both exported on every rank, with NVFP4 weights and an FP8 KV cache. Their config.json carries no base_model_ep_plan, as the source's has none, so no torch_a2a styles ship.

    Arm GPUs (nodes) Layout Result Wall time Shards
    A 4 (1) EP4, edp = 1 passed 8:17 4
    B 8 (2) EP4, edp = 2 passed 6:07 8
  • Rebased onto main 059588b, 2026-10-30, on [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's rebased head, without conflicts; this PR's own diff is unchanged. The GPU run above was at an earlier head with the same own diff. Pre-commit, mypy included, passes against main. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (857) and 5.5.4 (822 passed, 35 skipped), and without transformers the unit tests pass or skip (451 passed, 104 skipped), where import modelopt.torch.export does not import the writer. GPU tests not run here; CI defers them while this PR targets another PR's branch.

  • Rebased onto main d464d30, 2026-10-30, on [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's rebased head, without conflicts; this PR's own diff is unchanged, and the GPU run above, at an earlier head, had the same own diff. Pre-commit, mypy included, passes against main. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (872) and 5.5.4 (837 passed, 35 skipped), and without transformers the unit tests pass or skip (466 passed, 104 skipped), where import modelopt.torch.export does not import the writer. GPU tests not run here; CI defers them while this PR targets another PR's branch.

  • Text fixes, 2026-10-31, comments and a docstring only: [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 reworded the export comment this PR replaces, so this PR was restacked on it, keeping its own block; and _process_hf_model_for_export's docstring and the call-site comment no longer name a consolidation step. Pre-commit, mypy included, passes against main, and on CPU the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (872).

  • Rebased onto main e352cac, 2026-10-31, on [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752's rebased head, without conflicts; this PR's own diff is unchanged, and the GPU run above, at an earlier head, had the same own diff. Pre-commit, mypy included, passes against main. On CPU, the export unit tests, test_dataset_utils.py, test_distributed.py, test_moe_a2a.py and hf_ptq's argument tests pass on transformers 5.18 (872) and 5.5.4 (837 passed, 35 skipped), and without transformers the unit tests pass or skip (466 passed, 104 skipped), where import modelopt.torch.export does not import the writer. GPU tests not run here; CI runs them once this PR targets main.

Before your PR is "Ready for review"

  • Is this change backward compatible?: ✅ Only expert-parallel models, which could not be exported before, take the new route. A sharded export now raises on every rank when any rank's export fails, where it used to return on the ranks that succeeded.
  • 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?: ✅ on 2026-10-27; ⏳ a re-review of the code this PR took over from [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 on 2026-10-29 was requested.

Additional Information

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

Summary by CodeRabbit

  • New Features

    • Expert-parallel models can be exported without gathering weights, with ranks writing their assigned experts.
    • Export supports global and local quantization metadata and additional checkpoint tensor formats during conversion.
    • Expert-parallel export is enabled by default; FSDP2-only models can opt in to no-gather export.
  • Bug Fixes

    • Export detects unsupported sharding and inconsistent expert quantization, and coordinates failures across ranks.
    • Checkpoints preserve the source distributed configuration without changing the live model’s configuration.
    • Reversed quantization conversions preserve unrelated weights and restore expert names after dequantization or decompression.

@copy-pr-bot

copy-pr-bot Bot commented Oct 12, 2026

Copy link
Copy Markdown

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

coderabbitai Bot commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

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.

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: 6d702a3a-f127-44e4-ab2a-8417b6e31cf5
📥 Commits

Reviewing files that changed from the base of the PR and between dd732b4 and c2863a5.

📒 Files selected for processing: 6
  • modelopt/torch/export/quant_aware_conversion.py
  • modelopt/torch/export/unified_export_hf.py
  • modelopt/torch/utils/distributed.py
  • tests/unit/torch/export/test_quant_aware_conversion.py
  • tests/unit/torch/export/test_unified_export_hf.py
  • tests/unit/torch/utils/test_distributed_run.py

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

Walkthrough

HF export now handles expert-parallel sharding with coordinated no-gather writes. Quantization-aware reverse conversion handles load-time checkpoint rewrites and companion tensor patterns.

Changes: Expert-Parallel HF Export

Layer / File(s) Summary
Expert-parallel state and indexing
modelopt/torch/utils/distributed.py, modelopt/torch/export/distributed_writer.py, modelopt/torch/export/moe_utils.py, tests/unit/torch/utils/test_distributed_run.py, tests/unit/torch/export/test_distributed_writer.py
The writer obtains expert-parallel coordinates from routed experts and raises an error if the state is missing. Fused-expert quantizer lookup supports global-sized and EP-local lists.
HF shard selection and configuration
modelopt/torch/export/unified_export_hf.py, tests/unit/torch/export/test_unified_export_hf.py
HF export distinguishes FSDP2 sharding from expert parallelism, selects the no-gather writer for EP, rejects unsupported hybrid plans, and restores source distributed configuration values during serialization.
Distributed preparation and completion
modelopt/torch/export/unified_export_hf.py, tests/unit/torch/export/test_unified_export_hf.py
Export coordinates model preparation, fused-expert quantization checks, and completion across ranks. Processing and writing failures are reported across ranks.

Quantization-Aware Conversion Reversal

Layer / File(s) Summary
Reverse rules for rewritten checkpoint patterns
modelopt/torch/export/quant_aware_conversion.py, tests/unit/torch/export/test_quant_aware_conversion.py
Reverse rules recognize load-time quantization sibling suffixes, skip supported same-name codec conversions, and deduplicate module leaves for expert and dense conversion patterns.

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant HFExport as HF export
  participant AllReduce as all_reduce
  participant NoGatherWriter as no-gather writer
  HFExport->>AllReduce: Coordinate model preparation and quantization checks
  AllReduce-->>HFExport: Return collective result
  HFExport->>NoGatherWriter: Write sharded weights
  NoGatherWriter-->>HFExport: Return or report export failure
  HFExport->>AllReduce: Coordinate export completion
  AllReduce-->>HFExport: Return collective result
Loading
Merge Risk: Minimal · up to c2863

This change routes expert-parallel exports through the no-gather writer and coordinates failures across ranks. No concrete merge-blocking defect remains in the reviewed changes. Normal CI, including the GPU tests, should still run before merge.

Pre-merge checks: Passed 5
Check name Status Explanation
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns Passed PASS. The authoritative diff changes only five modelopt Python modules and tests; it changes no examples, pyproject.toml, or requirements files. Added package/example lines contain none of the prohibi…
Title check Passed The title names the main change: routing an expert-parallel Hugging Face model through the no-gather writer during export. The raw summary confirms this: expert-parallel exports use the no-gather writ…
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Full details: Security Anti-Patterns

Explanation

PASS. The authoritative diff changes only five modelopt Python modules and tests; it changes no examples, pyproject.toml, or requirements files. Added package/example lines contain none of the prohibited patterns: unsafe torch.load or numpy.load settings, hardcoded trust_remote_code=True, external-input eval/exec, or # nosec. No dependency is added.

Full details: Title check

Explanation

The title names the main change: routing an expert-parallel Hugging Face model through the no-gather writer during export. The raw summary confirms this: expert-parallel exports use the no-gather writer, and the objectives state the same aim. The title also carries the ticket ID and does not include file lists, emojis, or vague terms. It omits the supporting changes (reverse conversion, quantization agreement, rank coordination), which is expected for a short title.

Finishing Touches
Generate docstrings
  • Commit to this branch
  • Create a new PR
🧪 Generate unit tests (beta)
  • Commit to this branch
  • Create a new PR
Fix failing CI checks
  • Commit to this branch
  • Create a new PR
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

@shengliangxu
shengliangxu added this pull request to stack #2755 October 12, 2026 00:57
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from f3a1e12 to 4810e33 Compare October 12, 2026 01:11
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch 2 times, most recently from 4e11e09 to be3cc39 Compare October 12, 2026 02:44
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from 4810e33 to 72eab4a Compare October 12, 2026 02:44
@github-actions

github-actions Bot commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2753/

Built to branch gh-pages at 2026-10-26 23:48 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test be3cc39

@codecov

codecov Bot commented Oct 12, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.33333% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 62.36%. Comparing base (e79c725) to head (2ec61a3).

Files with missing lines Patch % Lines
modelopt/torch/export/distributed_writer.py 50.00% 4 Missing ⚠️
modelopt/torch/export/quant_aware_conversion.py 88.46% 3 Missing ⚠️
modelopt/torch/export/unified_export_hf.py 97.89% 2 Missing ⚠️
Additional details and impacted files
@@                           Coverage Diff                            @@
##           shengliangx/hfep-12c-nogather-export    #2753      +/-   ##
========================================================================
+ Coverage                                 62.30%   62.36%   +0.06%     
========================================================================
  Files                                       656      656              
  Lines                                     73764    73884     +120     
========================================================================
+ Hits                                      45955    46080     +125     
+ Misses                                    27809    27804       -5     
Flag Coverage Δ
unit 62.36% <93.33%> (+0.06%) ⬆️

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
shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch from be3cc39 to bdf08a2 Compare October 12, 2026 08:39
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from 72eab4a to a30a335 Compare October 12, 2026 08:39
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test bdf08a2

@shengliangxu shengliangxu changed the title [18/19] Export an expert-parallel model through the no-gather writer [OMNIML-6084] [18/19] Export an expert-parallel model through the no-gather writer Oct 12, 2026
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from a30a335 to 751b1d1 Compare October 13, 2026 00:53
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch from bdf08a2 to 136543b Compare October 13, 2026 00:53
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from 751b1d1 to 717b73e Compare October 13, 2026 01:32
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch from 136543b to 405ef82 Compare October 13, 2026 01:32
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 405ef82

@shengliangxu
shengliangxu removed this pull request from stack #2755 October 13, 2026 17:42
@shengliangxu
shengliangxu added this pull request to stack #2775 October 13, 2026 17:43
@shengliangxu shengliangxu changed the title [OMNIML-6084] [18/19] Export an expert-parallel model through the no-gather writer [OMNIML-6084] [18/20] Export an expert-parallel model through the no-gather writer Oct 13, 2026
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch from 405ef82 to 6fdea4b Compare October 13, 2026 19:38
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from 717b73e to 446370c Compare October 13, 2026 19:38
@shengliangxu
shengliangxu removed this pull request from stack #2775 October 13, 2026 19:38
@shengliangxu
shengliangxu added this pull request to stack #2781 October 13, 2026 19:40
@shengliangxu shengliangxu changed the title [OMNIML-6084] [18/20] Export an expert-parallel model through the no-gather writer [OMNIML-6084] [20/21] Export an expert-parallel model through the no-gather writer Oct 13, 2026
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-12c-nogather-export branch from 446370c to 755ebf2 Compare October 13, 2026 21:12
@shengliangxu
shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch from 6fdea4b to 2b63470 Compare October 13, 2026 21:12
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 2b63470

CI's partial-install torch lane runs tests/unit/torch without the hf extras. These two tests import transformers inside the test body, so they failed there instead of skipping like the rest of the file.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A barrier alone let the ranks whose writes succeeded return, and the run carried on as if the
checkpoint were whole while another rank's export raised. The sharded export now ends by
gathering every rank's outcome (all_reduce), so every rank raises with the failing rank's error.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…on plan

_hf_shard_flags reports FSDP2 and expert-parallel sharding from the model's recorded plan, rather
than from an FSDP2 module or a seeded expert-parallel group, so the export's route, its single
writer of the carried weights and its fail-together finish all follow the plan. The writer takes
its EP coordinate from the experts' group only when the plan is expert-parallel, and raises if
the experts carry none: without it every EP rank would write experts 0..n-1 over the others'.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Each of its three cases spawned two ranks, about 14 s of the linux unit job apiece.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With fewer expert-parallel ranks than the world, FSDP2 shards the dense part
within each expert-parallel group and replicates it across the groups. The
no-gather writer, which every expert-parallel export takes, writes a dense tensor
only from a dim-0 shard over the whole world, so it refused that layout -- after
the export had folded the quantizers. The export now refuses the plan before it
folds anything and names the knobs that avoid it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… failed

Each rank folds only its own shard, so processing can fail on one rank alone,
for instance when that rank's experts mix quantized and unquantized ones. That
rank went on to the closing export_ok consensus while its peers started the
writer's collectives, and the two did not match. Processing now ends in its own
consensus, so every rank raises before any of them writes.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…rallel export

Each rank holds only its own experts' weight quantizers, so the keep-fused
fold's all-or-none check saw only that rank's share. A rank whose experts were
all unquantized and one whose experts were all quantized both passed, folded
differently, and only the quantized one entered the writer's input-scale
all-reduce, out of step with its peer's ownership gather. An expert-parallel
export now gathers every rank's per-projection state before anything is folded,
and every rank refuses a projection whose experts are not all quantized or all
unquantized.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The two-rank test gave both projections the same states and never had every rank
mixed, so an agreement that read only the first projection, or accepted a
projection mixed on every rank, still passed. Both cases are now refused before
processing, as the agreement intends.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The tests of the closing consensus called _finish_export_together
directly, so they re-tested all_reduce: had export_hf_checkpoint gone
back to a plain barrier, every test would still have passed. One pooled
export now fails on rank 0 alone, after every rank wrote its shards, and
both ranks must raise with its error. Its collectives run on a group
that gives up in seconds, so a rank left waiting fails the test instead
of hanging it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…red ones

Merged:
- test_a_dequantizing_fp8_load_still_reverses_the_other_conversions becomes
  the fp8-dequantize case of test_experts_revert_through_converters_loading_rewrote,
  as a real dequantizing load leaves both converters: the case appends the
  in-place Fp8Dequantize converter and checks that a bf16 weight is not
  re-encoded. Aborting the reverse on that converter, and on transformers 5.5
  not recognizing its identity reverse, still fail the case.
- test_a_model_without_expert_parallel_experts_has_no_expert_state folds into
  test_the_routed_experts_state_is_found_by_its_expert_parallel_group.

Dropped, each caught by a remaining test when its code is broken:
- test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config:
  returning no quant config, folding nothing, or never keeping experts fused
  fails test_nogather_export_writes_what_the_single_process_export_writes.
- test_the_export_takes_how_the_model_is_sharded_from_its_plan: dropping either
  flag, swapping them, or reading a pipeline's plan fails the after-the-write
  failure test, the replicated-dense refusal or
  test_export_diffusers_pipeline_config.
- The one-rank-mixed case of the expert-quantization agreement: every mutation
  of the agreement it caught, the ranks-disagree and every-rank-mixed cases
  catch too.

The agreement cases now name the refused projection instead of pinning each
rank's state in the message, and the processing-failure job checks that rank 0's
error names rank 1. Both jobs, and the after-the-write job in place of its own
group, run their collectives inside bounded_collectives, and the
processing-failure job uses try/except rather than pytest.raises.

Kept, since nothing else fails when their code is broken: the agreement's
down-proj-disagrees case (an agreement that reads only the first projection)
and its all-unquantized case (one that refuses an agreed unquantized
projection).

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_dense_split_names_each_part_once_through_load_siblings imported
_dense_split_rule inside the test, next to the file's module-level
quant_aware_conversion imports. Its transformers imports stay local, with the
note the file already uses: CI's transformers-free lane collects this file.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…cope

The expert-parallel export tests imported mock, the distributed writer,
unified_export_hf and the distribution-plan helpers inside each test. Move
those to the module's import block; the transformers imports stay local, with
the file's existing note, since the transformers-free lane collects it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…nks agree

The expert-parallel export's agreement step read down_proj_weight_quantizers unguarded, while every
other walk of the pair treats the second projection as optional: a layout that registers only the
first projection's quantizers raised AttributeError inside the collective, on every rank.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its only caller is the no-gather writer, on a sharded model, which has already decided the route; a
second read of the plan that disagreed would have folded the experts split while the writer
expected them fused.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…place

They are a constant of the installed transformers, and were looked up again for every converter.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The no-gather export now saves the ModelOpt state from rank 0, as the streaming export does. An
expert-parallel model always takes that writer, and there rank 0 holds only its own experts' state,
so the save would write an incomplete one. The export refuses it up front, before any collective.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
transformers 5.5 imports triton atop its FP8 integration, so without
triton that import fails, and with it the shared import of
core_model_loading: the identity op 5.5 and 5.6 reverse Fp8Dequantize
into was lost too, and the reverse of a dequantizing load aborted, the
MiniMax-M2.7 failure. Each source of codec ops is now imported on its
own.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…rupted in

The closing consensus learned of a rank's failure only through
`except Exception`, so a SystemExit or KeyboardInterrupt on one rank
reached it as a success: the other ranks returned as if the checkpoint
were whole while that rank's export had stopped. The export now hands
any BaseException to the consensus, which, as all_reduce does, exits
every rank with an exit's code and fails them on anything else.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The writer's refusal of an expert-parallel model with no routed expert's
ParallelState, and expert_parallel_state's docstring, named
load_and_prepare_ep, which is not below this change. They now say what
such a load does instead: it seeds each routed expert's ParallelState.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… at the top

The writer read the expert count off the config with its own lookup,
which missed num_routed_experts and the llm_config a NemotronH-Omni
config nests its text model under.
modelopt.torch.utils.moe._num_experts_from_config reads every spelling
and both nestings, and returns 0 for a model that is not an MoE, as the
lookup did, so the writer uses it.

distribution_plan and expert_parallel_state were imported inside the
function. Nothing forces that: the writer already imports
modelopt.torch.utils.distributed at module scope, and moe.py imports only
torch. Both now come from the top.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_write_base_config asked the model whether it can generate before it
wrote generation_config.json. It now writes the non-weight files through
save_non_weight_artifacts, which does not ask, so the test's stub model
drops the answer.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The no-gather writer's local split now numbers experts by their FSDP2
dim-0 offset alone. Under expert parallelism transformers slices each
expert parameter to the rank's group, so that offset is group-relative:
the split takes the rank's expert-parallel coordinate and the config's
expert count again, adds the group's base when the tensor holds one
group's experts, and keeps the test of the numbering under each layout.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The writer reads each rank's expert-parallel coordinate and expert
replica index off the routed experts' ParallelState. A ParallelState
left with its default data-parallel group holds the whole world there:
taken as the replica index, it would leave every rank but global rank 0
out of the expert write, and the checkpoint would miss experts with no
error. The writer now refuses groups whose sizes are not the plan's
ep_size and edp, on every rank alike, as the sizes are the same on every
rank. A two-rank test seeds the default group and the plan's own.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ed weights

A sharded export reads the carried tensors off the source on rank 0
alone, and its peers read only their names. That read sat ahead of every
collective the export joins, so a failure there left rank 0 raising out
of the export while its peers waited in the first collective for it.

The read now runs inside an all_reduce for a sharded export, so a
failure on rank 0 raises on every rank, before any of them goes on. The
names every rank seeds get_quant_config with are rank 0's, the same on
every rank. A non-sharded export reads as before, with no collective.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…dule

A converter whose reverse only re-encodes a weight's quantization is
skipped, and its reverse rules with it, because it renames nothing. The
check compared only the first source and target, so a converter with a
later target naming another module would have been skipped too, leaving
that rename out of the exported names with no error. Every pattern must
now name a tensor of the same module, the weight or a scale the load
folded into it; another converter is reversed or refused, as before.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…k only when sharded

Only rank 0's list of carried names is used, so its peers no longer read
the source for names they would throw away: broadcast_from runs the read
on rank 0 alone, sends the list, and raises on every rank when that read
fails, as the all-reduce did.

The traceback print before the closing consensus is for a sharded export,
where the other ranks wait in it. A single process prints the traceback
once again when the error reaches the top, so it no longer prints it here.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The config check in _source_distributed_config_attrs imported PretrainedConfig inside the function, with no stated reason. transformers is optional for this module, and the module already imports it once, guarded, at the top. PretrainedConfig joins that import, and the check reads _TRANSFORMERS_CONFIG_TYPES, which is empty without transformers, as _TRANSFORMERS_MODEL_TYPES is.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The no-gather export has no consolidation step: each rank's writer splits
its own fused experts into per-expert tensors. #2752 fixed its own copy of
this wording; the docstring and comment this PR adds said the same. Comments
and a docstring only.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 8d71098

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants