Repository navigation
[OMNIML-6084] Export an expert-parallel model through the no-gather writer - #2753
Open
shengliangxu wants to merge 37 commits into
Open
shengliangxu wants to merge 37 commits into
shengliangxu wants to merge 37 commits into
Conversation
|
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. |
Contributor
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 12, 2026 01:11
f3a1e12 to
4810e33
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-13-ep-export
branch
2 times, most recently
from
October 12, 2026 02:44
4e11e09 to
be3cc39
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 12, 2026 02:44
4810e33 to
72eab4a
Compare
Contributor
|
Collaborator
Author
|
/ok to test be3cc39 |
Codecov Report❌ Patch coverage is
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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
shengliangxu
force-pushed
the
shengliangx/hfep-13-ep-export
branch
from
October 12, 2026 08:39
be3cc39 to
bdf08a2
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 12, 2026 08:39
72eab4a to
a30a335
Compare
Collaborator
Author
|
/ok to test bdf08a2 |
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 13, 2026 00:53
a30a335 to
751b1d1
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-13-ep-export
branch
from
October 13, 2026 00:53
bdf08a2 to
136543b
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 13, 2026 01:32
751b1d1 to
717b73e
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-13-ep-export
branch
from
October 13, 2026 01:32
136543b to
405ef82
Compare
Collaborator
Author
|
/ok to test 405ef82 |
shengliangxu
force-pushed
the
shengliangx/hfep-13-ep-export
branch
from
October 13, 2026 19:38
405ef82 to
6fdea4b
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 13, 2026 19:38
717b73e to
446370c
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-12c-nogather-export
branch
from
October 13, 2026 21:12
446370c to
755ebf2
Compare
shengliangxu
force-pushed
the
shengliangx/hfep-13-ep-export
branch
from
October 13, 2026 21:12
6fdea4b to
2b63470
Compare
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>
Collaborator
Author
|
/ok to test 8d71098 |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_statetest moved totests/unit/torch/utils/test_distributed.py, which replacedtest_distributed_run.pyon main, and the writer's refusal and that function's docstring no longer name #2741'sload_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 andwrite_experts), which #2752 held at 0 and dropped, and_split_local_fused_module'sep_rank/total_experts, which #2752 also dropped. Review fixes since the first re-cut are listed at the end of Testing.Series layout
main([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 have merged)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_pretrainedsees a tensor-parallel model and refuses.export_hf_checkpointnow 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: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.config.json, and the closing barrier was skipped, so the job hung.distributed_config, and the source's ownbase_model_ep_planrather than thetorch_a2aone the loader swapped in, which only exists while ModelOpt is imported.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.Fp8Quantize), which the reverse could not do, so it gave up on the whole mapping, and MiniMax-M2.7 exportedmlp.experts.N.gate_projinstead ofblock_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.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: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, throughexport_hf_checkpointitself on two gloo ranks: rank 0's write ofconfig.jsonfails 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/uniton 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:
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_loadingare 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 aSystemExitorKeyboardInterrupton one rank fails every rank. The test'sSystemExitcase 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 theexpert_parallel_statetest moved totest_distributed.py. The writer's refusal of expert-parallel experts without aParallelStatesays 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.pyand hf_ptq's argument tests pass (801;test_export_diffusers.pyleft 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 readsnum_routed_expertsandllm_config), anddistribution_plan/expert_parallel_stateare imported at module scope._write_base_configis now a wrapper: it calls [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752'ssave_non_weight_artifactsinside 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.pyand 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_sizeandedp, 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.pyand 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), whereimport modelopt.torch.exportdoes 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 anddown_projhold 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.pyand 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), whereimport modelopt.torch.exportdoes 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.pyand 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), whereimport modelopt.torch.exportdoes not import the writer. GPU tests not run; CI runs them.Review fix, 2026-10-30 (meenchen's review bot):
PretrainedConfigmoves into the module's guarded transformers import at the top, and_source_distributed_config_attrschecks_TRANSFORMERS_CONFIG_TYPES, empty without transformers, as_TRANSFORMERS_MODEL_TYPESis. No behaviour change:test_distributed_export_writes_the_config_a_single_gpu_export_writes, whose plan sits in a VLM's nestedtext_config, passes, and fails with the config types emptied.test_unified_export_hf.pypasses 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 atc8ecddb3342. In the writer's files, this head differs from that one only by thePretrainedConfigimport move above. Qwen/Qwen3-30B-A3B, recipegeneral/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. Theirconfig.jsoncarries nobase_model_ep_plan, as the source's has none, so no torch_a2a styles ship.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.pyand 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), whereimport modelopt.torch.exportdoes 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.pyand 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), whereimport modelopt.torch.exportdoes 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.pyand 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.pyand 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), whereimport modelopt.torch.exportdoes not import the writer. GPU tests not run here; CI runs them once this PR targets main.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models; merges after #2752.
Summary by CodeRabbit
New Features
Bug Fixes