[OMNIML-6084] [17/21] Export an expert-parallel model through the no-gather writer by shengliangxu · Pull Request #2753 · NVIDIA/Model-Optimizer

shengliangxu · 2026-10-12T00:57:21Z

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 ([16/21]) unconditionally, since it has no fallback. It reads that the model is expert-parallel from the model's distribution plan ([4/21]), as it reads FSDP2 sharding ([16/21]). 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 ([9/21]), 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. The DTensor's global extent tells the two apart. (This was an IndexError on every rank exporting GLM-5.3-Flash at EP16.)
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 ([21/21] 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, [4/21]), so every rank raises with the failing rank's error.

Usage

None: export_hf_checkpoint(model, export_dir=...) on a model from load_and_prepare_ep, 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_run.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 reordered series rebased on main 264fd26: the full unit suite on transformers 5.18, and on 5.5 apart from the tests whose ONNX, TensorRT or diffusers packages that environment lacks; CI's transformers-free unit lane; pre-commit over the whole PR against main; and the Sphinx docs build. The same at every PR of the stack.
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 [21/21].

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 [21/21].
Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

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

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 · 2026-10-12T00:57:25Z

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 · 2026-10-12T00:57:31Z

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: c2863a5
Latest commit, not reviewed: 6922256

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.

Priority: Normal

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

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: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 74.51% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 51 functions across 9 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 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…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the primary change: exporting expert-parallel models through the no-gather writer.

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.

✨ Finishing Touches: 💡 2

📝 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

_{Comment @coderabbitai help to get the list of available commands.}

github-actions · 2026-10-12T02:51:25Z

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

shengliangxu · 2026-10-12T04:08:30Z

/ok to test be3cc39

codecov · 2026-10-12T04:08:51Z

Codecov Report

❌ Patch coverage is 92.25352% with 11 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (shengliangx/hfep-12c-nogather-export@920c4f7). Learn more about missing BASE report.

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/moe_utils.py	66.66%	2 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	97.91%	2 Missing ⚠️

Additional details and impacted files

@@                           Coverage Diff                           @@
##             shengliangx/hfep-12c-nogather-export    #2753   +/-   ##
=======================================================================
  Coverage                                        ?   62.01%           
=======================================================================
  Files                                           ?      656           
  Lines                                           ?    73553           
  Branches                                        ?        0           
=======================================================================
  Hits                                            ?    45613           
  Misses                                          ?    27940           
  Partials                                        ?        0

Flag	Coverage Δ
unit	`62.01% <92.25%> (?)`

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 · 2026-10-12T09:43:30Z

/ok to test bdf08a2

shengliangxu · 2026-10-13T02:37:59Z

/ok to test 405ef82

shengliangxu · 2026-10-13T21:32:23Z

/ok to test 2b63470

shengliangxu · 2026-10-23T07:22:15Z

/ok to test 951faa3

shengliangxu · 2026-10-23T08:05:01Z

/ok to test 0bb80b2

shengliangxu · 2026-10-23T17:07:08Z

/ok to test 6c7c6c0

shengliangxu · 2026-10-23T21:38:54Z

/ok to test 1f71c36

shengliangxu · 2026-10-23T22:34:04Z

/ok to test 777ff13

shengliangxu · 2026-10-23T23:54:37Z

/ok to test f02e2d7

shengliangxu · 2026-10-24T00:43:00Z

/ok to test 332db97

shengliangxu · 2026-10-24T01:24:25Z

/ok to test 7cb7af4

An EP-sharded 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 outright. Routing now sends an expert-parallel model to the no-gather writer unconditionally, since it has no fallback; a model sharded by FSDP2 alone keeps its opt-in. What expert parallelism changes in the writer and around it: - The global expert index. Under FSDP2 alone a per-expert key's index is the DTensor's own dim-0 offset, but under EP transformers has already sliced the parameter to this rank's group, so that offset is group-relative and every group would write experts 0..n-1 over each other -- a silently wrong checkpoint. The coordinate now comes from the routed experts' ParallelState (expert_parallel_state), the same state calibration reads. - Per-expert weight quantizers. The fold indexed them by the DTensor's global dim-0 offset, the FSDP2 convention where the quantizer list stays global; under EP the list holds only this rank's experts. What tells the two apart is the DTensor's global dim-0 extent -- equal to the list's length under FSDP2, larger under EP -- and an assert pins the local invariant. (Seen as an IndexError on every rank exporting GLM-5.3-Flash under EP16.) - The rank. It was get_rank() only under FSDP2 and 0 otherwise, so under EP every process took the rank-0 branch and rewrote config.json at once, and the closing barrier was skipped, so the job hung instead of exiting. - Carried weights. Under EP every other rank's experts read as unplaced on each rank, and every rank believed it was the single sender of the carried set; the writer's point-to-point protocol then deadlocked (GLM-5.3-Flash, EP16: 15 of 16 ranks parked in the barrier). Rank 0 alone writes them now. - transformers stamps a DistributedConfig on the config when sharding; it is not JSON-serializable and describes this run's mesh, so it is dropped first. - A rank whose export raises prints its traceback before the finally clause's barrier, where it would otherwise block on ranks that never fail and the job would hang with no error in the log. _hf_shard_flags derives (fsdp2_sharded, ep_sharded) once for all of these gates, kept separate because the FSDP2-dense + EP-expert hybrid is both at once. _process_hf_model_for_export names the in-place processing half the writer path runs: prepare, pack quantized weights with sharded experts left fused, return the quant config. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…json A distributed export wrote a config.json that differed from a single-GPU export of the same checkpoint in two keys: - base_model_ep_plan, which the EP loader swaps to the torch_a2a styles. Those styles exist only while ModelOpt is imported, so loading the export expert-parallel anywhere else fails on an unknown style. transformers drops the tp and pp plans when it serializes a config, but not the EP plan. - distributed_config: null. The writer nulled the DistributedConfig that transformers stamps on the config (it is not JSON-serializable), and the null was written out. Nulling it in place also left transformers' tp_plan property raising on the live model, since it reads distributed_config.enable_expert_parallel whenever the attribute exists. While config.json is written, each config -- the root and every nested one, where a VLM keeps its plan -- now holds the value its source config.json declares, or loses the attribute when the source has none. The run's values come back afterwards. Found comparing a Qwen3-30B-A3B export from the EP path with a single-GPU export of the same checkpoint. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Its docstring and test named the speculative-decoding exporters as a second caller; in this stack the no-gather writer is the only one. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Kimi-K3's expert-parallel prepare hook installs _checkpoint_conversion_mapping on the root and text configs, renaming the source's per-expert w1/w3/w2 so transformers can fuse them on load. Those regexes name the source's tensors, not the export's, and the distributed writer serialized them into config.json at both levels: the 16-node Kimi-K3 export carried them while its base_model_ep_plan and distributed_config were already gone. Restore them with the other run attributes, so a reader of the export never sees renames meant for a different checkpoint. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The export reverts in-memory names to the checkpoint's by replaying the conversions from_pretrained recorded. A load that rewrote a converter lists every checkpoint tensor it consumes as a source: a dequantizing FP8 load adds each weight's weight_scale_inv, and an on-the-fly compressed-tensors decompressor turns each .weight into weight_packed / weight_scale / weight_global_scale. The same expert leaf then repeats, the expert reverse bailed with an arity mismatch, and the whole reverse was skipped, leaving the export on in-memory names. On Kimi-K3 that wrote the routed experts as block_sparse_moe.experts.N.gate_proj/up_proj/down_proj, while the model's own modeling code builds them as w1/w3/w2 (gate_proj and friends belong to its dense MLP class), so a stock load matched none of the 989,184 expert tensors. A converter rewritten by an FP8 dequantizing load fails the same way. Read a module's leaf through the regex anchor and the compressed-tensors sibling suffixes, and keep each leaf once, in order. The dense split rule gets the same treatment, for a fused dense converter a dequantizing load rewrote. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

They were kept out of config.json because Kimi-K3's prepare hook installed _checkpoint_conversion_mapping on its configs. The hook registers its per-expert conversion with transformers' conversion registry instead now, and nothing in the expert-parallel path sets that attribute any more. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

A FineGrainedFP8 dequantize=True load adds a converter that decodes each weight in place, and its reverse re-encodes it (Fp8Quantize), which the quant-aware reverse does not support. It then gave up on the whole mapping, so MiniMax-M2.7 exported its routed experts under in-memory names (mlp.experts.N.gate_proj) instead of the checkpoint's (block_sparse_moe.experts.N.w1). Such a converter renames nothing, and the export writes ModelOpt's own quantized weights and scales rather than re-encoding the source's format, so it contributes no reverse rule now. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

transformers 5.5 and 5.6 reverse Fp8Dequantize into their identity op rather than Fp8Quantize, so on CI's minimum transformers the dequantizing converter still aborted the reverse. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

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>

shengliangxu · 2026-10-24T05:15:01Z

/ok to test 6922256

shengliangxu added this pull request to stack #2755 October 12, 2026 00:57

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

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 force-pushed the shengliangx/hfep-13-ep-export branch from bdf08a2 to 136543b Compare October 13, 2026 00:53

shengliangxu force-pushed the shengliangx/hfep-13-ep-export branch from 136543b to 405ef82 Compare October 13, 2026 01:32

shengliangxu removed this pull request from stack #2755 October 13, 2026 17:42

shengliangxu added this pull request to stack #2775 October 13, 2026 17:43

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 force-pushed the shengliangx/hfep-13-ep-export branch from 405ef82 to 6fdea4b Compare October 13, 2026 19:38

shengliangxu removed this pull request from stack #2775 October 13, 2026 19:38

shengliangxu added this pull request to stack #2781 October 13, 2026 19:40

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 force-pushed the shengliangx/hfep-13-ep-export branch from 6fdea4b to 2b63470 Compare October 13, 2026 21:12

shengliangxu added 21 commits October 24, 2026 05:11

Describe _process_hf_model_for_export's one caller

551f002

Its docstring and test named the speculative-decoding exporters as a second caller; in this stack the no-gather writer is the only one. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Keep two blank lines after expert_parallel_state

de37197

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

Run the every-rank-fails export test on the shared gloo workers

1d68642

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>

Conversation

shengliangxu commented Oct 12, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Oct 12, 2026

Uh oh!

coderabbitai Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Sequence Diagram(s)

Uh oh!

github-actions Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

codecov Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

shengliangxu commented Oct 12, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 13, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 23, 2026

Uh oh!

shengliangxu commented Oct 24, 2026

Uh oh!

shengliangxu commented Oct 24, 2026

Uh oh!

shengliangxu commented Oct 24, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

shengliangxu commented Oct 12, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading

github-actions Bot commented Oct 12, 2026 •

edited

Loading

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

codecov Bot commented Oct 12, 2026 •

edited

Loading