[shengliangxu]

Note

Re-cut again on 2026-10-29, so more of the series merges in parallel: this PR's base branch was rebuilt from the current #2751 and #2750 (both now on main; #3099, [1/2] of #2750's split, merged) and #2740 (on #3074). #2779 left it: this PR needed it only for the plan that hf_ptq's --use_fsdp2 load records, and that recording now lands here. The series first became a tree on 2026-10-27, from main 1521053 (this PR was [16/21]); it was rebased onto main 7c3c868 later that day, after #2739 merged, and onto main f03651e that night.

Base: shengliangx/hfep-base-12c-nogather-export, a branch (not a PR) that merges main and #2740, the PR this one still builds on (#2750 and #2751 have merged); #2740 brings #3074 below it. The diff below is this PR's alone (420 source lines), and it merges after all of them.

Re-cut: this PR's diff gained what the PRs below it handed over, listed in the last Testing line: _split_local_fused_module, the per-rank expert split this writer calls, with its two tests (from #2750's split); hf_ptq's --use_fsdp2 plan recording, with its README notes, changelog entry and test (from #2740); no out-of-memory retry in the calibration loop for a model over several ranks (asked for in #2740's review); and the writer's module-scope imports from #2751. 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)
  • [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm
    • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
    • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
  • [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 load routing
    • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 expert-parallel loader (on [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
      • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
        • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
      • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
        • [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 expert parallelism in the model specs
          • [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 Kimi-K2.5 and Kimi-VL adapters (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747)
          • [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747)
          • [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)
      • [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 on-the-fly decompression
  • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6071] Calibrate a distributed model on per-rank data #2740) (this PR)
    • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export

What does this PR do?

Type of change: New feature, with a bug fix for models built from a config

distributed_save_hf_checkpoint writes an already-processed sharded model as HF safetensors with no rank-0 gather:

• fused experts are split shard-locally ([OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750);
• every other tensor moves whole to one owner rank, which runs the reverse conversion ([OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751);
• each rank writes its own shards.

No rank ever holds the whole model. On top of its own shard, which the model keeps, a rank holds what it writes; only an inference tensor is copied up front, so an export under no_grad, as hf_ptq runs it, copies nothing. A model sharded by FSDP2 alone opts in with MODELOPT_FSDP2_EXPORT=nogather; the expert-parallel export routes to it unconditionally in #2753.

Two gloo ranks exporting a 223 MiB bf16 Llama, the writer's peak RSS above each rank's 112 MiB shard:

Writer	Peak RSS above the shard, per rank
copying every DTensor up front (previous revision)	278 / 246 MiB
copying only inference tensors	224 / 217 MiB
copying only inference tensors, with #2751's send-copy fix (this revision)	162 / 166 MiB

Why this doesn't build on torch.distributed.checkpoint's HuggingFaceStorageWriter: that writer saves each state-dict key as it is, and the checkpoint needs other keys, made from whole tensors:

• Per-expert keys. The model holds a layer's experts fused, as mlp.experts.gate_up_proj of shape [experts, 2 × intermediate, hidden], while an HF checkpoint names each expert's projections (mlp.experts.3.gate_proj.weight). DCP can map a key to a file (fqn_to_index_mapping) but cannot write one key as many, so the split comes first. It is shard-local ([OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750) and leaves each rank's experts as plain tensors, whole on that rank, which leaves a sharded writer nothing to do for them.
• The reverse conversion needs whole tensors. transformers' conversion ops take whole tensors and know no placement, and some reverse merges concatenate along dim 0, the dim FSDP2 shards. hrm_text, for one, loads attn.gqkv_proj.weight with Chunk(dim=0) into gate, q, k and v, so the export concatenates the four back. Concatenating each rank's row shards gives [gate_r; q_r; k_r; v_r], which is no row range of the whole. So each tensor has to be whole somewhere before the conversion runs: this writer makes it whole on one owner rank, rather than everything on rank 0.
• Files a loader reads. With save_distributed=True and no consolidation, a DTensor's chunks land in several files under one key, placed by offsets in DCP's own metadata, and no model.safetensors.index.json is written: only DCP's reader reassembles them, not from_pretrained, vLLM or TensorRT-LLM. With enable_consolidation=True, the writer's finish runs consolidate_safetensors_files on the coordinator rank, which reads every sharded file back and writes the checkpoint again: every byte is written twice, the second time by rank 0 alone. With save_distributed=False, rank 0 saves a full state dict, the gather this path avoids.

Once every tensor is whole on exactly one rank, each rank writes its own set once with safetensors.save_file, and rank 0 unions the index.

export_hf_checkpoint on this path:

• packs quantized weights with keep_fused_experts;
• writes config.json and a sanitized generation_config.json on rank 0 through save_non_weight_artifacts, as the FSDP2 streaming writer does, then hf_quant_config.json. hf_quant_config.json names modules as the weights do: when any rank cannot reverse the conversion mapping, every rank's weights keep their in-memory names, and the quant config keeps them too (distributed_save_hf_checkpoint returns whether the names were reverted);
• copies the source's non-model files (tokenizer, processor, remote code, chat templates), which the writer never sees, also through save_non_weight_artifacts.

With save_modelopt_state, it saves the ModelOpt state from rank 0, as the FSDP2 streaming export does.

The export reads how the model is sharded from its distribution plan (#2736), not from an FSDP2 module in it: the FSDP2 routes, streaming or no-gather, and which rank reads the carried-over weights. hf_ptq's --use_fsdp2 load records its plan, DistributionPlan(world_size, dense="fsdp"), right after the load. A model sharded by hand, with fsdp2_wrap say, records it with set_distribution_plan; without it the export raises rather than treat it as whole. Refusing a model with no plan, rather than inferring DistributionPlan(world_size, dense="fsdp") from its FSDP2 modules, is deliberate: the export reads the plan through #2736's distribution_plan, which raises for such a model, because an FSDP2 module does not say how the model is distributed. A guessed plan would leave out the expert-parallel degree, which #2753 numbers each rank's experts by. A diffusers pipeline is not a module and is exported without a plan.

The calibration loop no longer retries an out-of-memory batch in halves for a model whose recorded plan spans several ranks, such as hf_ptq's --use_fsdp2 model now that it records one: its forward can hold collectives, so a rank that retried would run more forwards than its peers, which would wait in them. The error is raised instead, and the README tells --use_fsdp2 users to set --batch_size. Recording the plan also switches on #2740's refusal of calibration that keeps statistics per rank and its collective batch sizing for --use_fsdp2 on several GPUs, so their README notes and backward-breaking changelog entry come here too.

The KV-cache format goes through _get_kv_cache_postprocess_config, so a mixed-precision KV cache exports as on the other paths, and which writer runs is rank 0's MODELOPT_FSDP2_EXPORT.

Usage

MODELOPT_FSDP2_EXPORT=nogather torchrun --nproc_per_node 8 examples/hf_ptq/hf_ptq.py --use_fsdp2 ...

Testing

• tests/unit/torch/export/test_distributed_writer.py adds end-to-end cases on CPU under gloo:
  • One rank:
    • an FSDP2 Mixtral written byte-identical to save_pretrained, tied and untied;
    • inside inference mode;
    • extra_state_dict kept under its own names;
    • in each of these, every dense DTensor reaching the transfer as the model's own shard, not a copy (it fails with the unconditional copy, mutation-checked);
    • export_hf_checkpoint on the no-gather path, root-wrapped and layers-only: byte-identical to the single-process export, plus config files and the copied source sidecar, for FP8 and NVFP4;
    • a Qwen2-MoE's experts exported quantized, for FP8 and NVFP4, their inputs too, so the experts' input scales and their all-reduce run on CPU. Its shared expert is packed after the experts, so experts folded outside FSDP2's weight window ([OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750) would be written unquantized; this replaces [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's stand-in test of that window;
    • a quantized Mixtral whose reverse conversion fails, keeping its in-memory names in exclude_modules as in the weights;
    • save_modelopt_state saved from rank 0.
  • Two ranks:
    • an uneven Mixtral matching its source, with one continuous file numbering;
    • an FP8 dense model matching the single-process export;
    • extras handed to rank 0 alone, as the export entry hands them, written exactly once, and an extra colliding with a model weight raising on both ranks;
    • each tensor with exactly one owner, both ranks owning some;
    • divergent dense keys, sharding and duplicate names across ranks raising instead of hanging;
    • ranks that calibrated the experts' shared input quantizer differently, every expert written with the largest scale.
  • The Mixtral-names and two-rank input-scale cases each fail without their fix (mutation-checked).
• tests/unit/torch/export/test_unified_export_hf.py: a model built from a config, which main's Speed up CPU unit tests to stay under the 20 min CI limit #2823 gives no source checkpoint, copies nothing from the working directory at export.
• The plan reads: reading none, taking no FSDP2 route, or dropping the pipeline guard each fails a test, the last main's test_export_diffusers_pipeline_config on a real tiny pipeline. Only rank 0 reading the carried weights is not unit-tested, since it needs two FSDP2 ranks; the GPU carry-over tests cover it.
• Mutation-checked, across this PR and [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751: 76 of 77 single-point mutants killed. The survivor is a copy of plain inference tensors that nothing on CPU can reject.
• The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.
• New tests/gpu/torch/export/test_distributed_hf_export.py: PTQs a tiny Qwen3 (dense, with tied embeddings, and MoE) under FSDP2 with FP8 and NVFP4, opting each export into the writer, and checks the checkpoint against the safetensors headers: every routed expert exported quantized, and the tied head written once. Its earlier version passed on a 4-GPU node, 8 tests, including checks of extras, ties, refusals and unquantized layers that now run on CPU above, where every PR runs them. On 1 GPU, all skip. Its docstring no longer describes a DCP staging step this writer never had, and an assertion on a sharded/ staging directory, which could not fail, is dropped; not re-run on GPU after that change.
• The end-to-end cases fold into the hub-checkpoint test and the two-rank job, and tests a remaining test covers are dropped, this PR's test_export_diffusers.py test included. Each drop is checked by mutation.
• The single-process comparison has an FP8 and NVFP4 KV-cache case, and the two-rank export sets the writer knob on rank 0 only.
• 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.
• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.
• Review fixes, 2026-10-27: the ModelOpt state is saved from rank 0, as the streaming export does; the ranks agree on the expert input-scale keys before reducing them; the writer is told the model's merge groups. 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: MODELOPT_FSDP2_EXPORT takes nogather or streaming (the default) and refuses any other value on every rank.
• Review suggestion, 2026-10-27 night: the no-gather arm discards the tied-weight map, with a comment that the writer builds its own.
• Re-cut on 2026-10-29 (see the note at the top), onto the base rebuilt from [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 bb9dd1f, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 830d704 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 54f0c1f. The rebase's conflicts were all in tests and the changelog: imports, the changelog's 0.49.0 sections, [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751's revised test helper (its signature kept; the two _revert_whole_sd tests this PR drops still go, since the end-to-end tests fail when that code breaks), and [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's stand-in window test, dropped with the imports only it used (its window test now lives in [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's test_sharded_weight_packing.py). Added on top: _split_local_fused_module with test_split_local_numbers_experts_globally and each rank's own split checked in [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's two-rank fold test; hf_ptq's plan recording with test_the_fsdp2_load_records_the_plan_it_loaded_by; the no-retry with a test that fails without it (both a loop that always splits and one that splits for a one-rank plan fail it); and the module-scope imports, with the test that fails the reverse conversion now patching the name the writer calls (it failed on the new base, and passes). Pre-commit, mypy included, passes against the base 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 (782; test_export_diffusers.py left out: its 7 failures here are main's too). At [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753's head, which carries this PR, they also pass on transformers 5.5.4. A no-gather export on one gloo rank writes INT4-AWQ, MXFP4, NVFP4 and FP8 Llama and Qwen2-MoE checkpoints byte-identical to the single-process export, so [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099's refusal to pack a DTensor shard in a 4-bit packer is not reached: the fold packs each rank's experts as plain tensors, and the dense weights are packed inside FSDP2 weight windows. GPU tests not re-run; CI runs them.
• Review fix, 2026-10-29: the GPU carry-over export tests and the GPU loader tests record the plan of the model they shard by hand, as the export now requires. Their lines are the ones [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 adds to the same tests, and that PR is no longer below this one. The changelog entry names parallel_load_and_prepare_fsdp2 next to fsdp2_wrap as recording no plan, and the comment above the held FSDP2 weight windows says the no-gather export sets keep_fused_experts. Pre-commit, mypy included, passes against the base and against main. The two GPU files collect on CPU. On one gloo rank, the default FSDP2 export of a model with a recorded plan is byte-identical to the single-process export (FP8 and NVFP4, Llama and Qwen2-MoE). The GPU suites were not run.
• Review fixes, 2026-10-29 night (meenchen's review): the writer drops the expert-parallel scaffolding it never set (ep_rank/dp_idx held at 0, the expert count); [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 adds it back with what sets it. The writer imports its helpers at module scope, and the export's import of the writer stays local, with a comment naming the circular import a module-scope one causes. Rank 0 writes the non-weight files through save_non_weight_artifacts, as the FSDP2 streaming writer does: _write_base_config and the special case in the final copy are gone. Quantized Llama and Qwen2-MoE no-gather exports write byte-identical config, generation config, quant config, index and copied files before and after. 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 (829) and 5.5.4 (794 passed, 35 skipped); without transformers, the unit tests pass or skip (438 passed, 93 skipped) and import modelopt.torch.export does not import the writer. Pre-commit, mypy included, passes against the base and against main. GPU tests not run; CI runs them.
• Review suggestion, 2026-10-30: _split_local_fused_module numbers experts by the FSDP2 offset alone. Its expert-parallel parameters and their numbering test move to [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753, which sets them. 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 (824) and 5.5.4 (789 passed, 35 skipped), and without transformers (438 passed, 93 skipped). Pre-commit, mypy included, passes against the base and against main.
• Rebased onto main 99c3834, 2026-10-30, after [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 merged: the base branch now merges [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 only. This PR's own diff is unchanged except for one docstring, where split_fused_experts_state_dict's note on memory and expert numbering now merges [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750's final wording with this PR's. [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 _split_local_fused_module, which passes each rank's local shards of both, and FSDP2's dim-0 sharding gives them equal counts; the two-rank fold test, three experts over two ranks included, passes. 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 (825) and 5.5.4 (790 passed, 35 skipped), and without transformers the unit tests pass or skip (438 passed, 93 skipped). Pre-commit, mypy included, passes 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 (and [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122, below [OMNIML-6071] Calibrate a distributed model on per-rank data #2740): the base branch now merges main and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740. This PR's own diff is unchanged. Its one conflict was changelog context: main's new Qwen3.8 entry now comes before this PR's no-gather entry, and both stay. 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 (827) and 5.5.4 (792 passed, 35 skipped), and without transformers the unit tests pass or skip (438 passed, 95 skipped). GPU tests not run; CI runs them.
• Rebased onto main 059588b, 2026-10-30 (after Skip long-lived objects in the per-step gc.collect() of QAT and QAD #3104, [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128, [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 and ci: establish a shared review baseline and adopt it in Claude reviews #3149 merged), on [OMNIML-6071] Calibrate a distributed model on per-rank data #2740's rebased head: the base branch again merges main and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740. This PR's own diff is unchanged, plus one commit: [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128 now refuses AWQ with --use_fsdp2 on more than one rank before the model loads, so this PR's README note lists only the other calibrations it refuses when quantization starts (SVDQuant, local Hessian, imatrix) and points to [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's AWQ note. The restack's conflicts were in the changelog and the README only, where main's new entry and notes sit next to this PR's; all are kept. 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 (828) and 5.5.4 (793 passed, 35 skipped), and without transformers the unit tests pass or skip (438 passed, 95 skipped). GPU tests not run; CI defers them while this PR targets another PR's branch.
• Rebased onto main d464d30, 2026-10-30, after [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged (and fix(export): support stateless constant KV scales in functional capture #3121), on [OMNIML-6071] Calibrate a distributed model on per-rank data #2740's head now on main: without conflicts, and 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 (843) and 5.5.4 (808 passed, 35 skipped), and without transformers the unit tests pass or skip (453 passed, 95 skipped). GPU tests not run; CI defers them while this PR targets another PR's branch.
• Review text fixes, 2026-10-31 (meenchen's review bot), comments and a docstring only: the two-rank test's comment and docstring, and the writer's comment on the tensors the model does not own, say rank 0 alone is handed them, as the list above now does; the export's comment on packing no longer names a consolidation step or main's history; and the writer's step (2) comment no longer names an expert-parallel layout. Pre-commit, mypy included, passes against main, and test_distributed_writer.py and test_unified_export_hf.py pass (123). The GPU suites are pending: tests/gpu/torch/export/test_distributed_hf_export.py, tests/gpu/torch/export/test_export_carry_over.py and tests/gpu/torch/utils/test_model_load_utils.py have not run at this head; CI runs them once this PR targets main.
• Rebased onto main e352cac, 2026-10-31 (after Run fused-MoE expert GPTQ updates on all of this process's GPUs #3002, Collect imatrix statistics with one fused float32 reduction #3044 and [OMNIML-6273] Let the Claude review bot submit its review again #3156 merged), on [OMNIML-6071] Calibrate a distributed model on per-rank data #2740'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 (843) and 5.5.4 (808 passed, 35 skipped), and without transformers the unit tests pass or skip (453 passed, 95 skipped). The GPU suites are still pending: CI runs them once this PR targets main.

Before your PR is "Ready for review"

• Is this change backward compatible?: ⚠️ The no-gather writer is opt-in for FSDP2. A model FSDP2-sharded by hand, fsdp2_wrap included, must record its plan with set_distribution_plan before export_hf_checkpoint; without it the export raises. hf_ptq.py --use_fsdp2 on several GPUs now refuses calibration that keeps its statistics on each rank, and raises on a calibration batch that runs out of GPU memory instead of retrying it smaller.
• 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?: ✅ The no-gather FSDP2 export opt-in, and, under Backward Breaking Changes, that the export needs a distributed model's plan recorded and what hf_ptq.py --use_fsdp2 on several GPUs now refuses.
• Did you get Claude approval on this PR?: ✅ Approved on 2026-10-29 after meenchen's round. Its one suggestion was taken in a later round.

Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models; merges after #2751, #2750 and #2740.

Summary by CodeRabbit

• New Features
  • Added an opt-in no-gather export for FSDP2-sharded models. With MODELOPT_FSDP2_EXPORT=nogather, ranks write checkpoint shards in parallel without gathering the full model on one host; MoE experts remain sharded.
  • Distributed exports include non-model files from the source checkpoint when available.
• Bug Fixes
  • Export and batch-size checks now use the distribution plan recorded during model loading. Manually distributed models need a plan recorded before calibration or export; operations raise an error when a required plan is missing.

[copy-pr-bot]

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

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 12267ec
• Latest commit, not reviewed: ac29323

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

Walkthrough

Adds an opt-in FSDP2 no-gather Hugging Face export path. The exporter uses the distribution plan and environment setting to select the path. Distributed ranks process and write checkpoint shards. Tests cover dense and MoE exports, quantization, metadata, and multi-rank writing.

Changes: FSDP2 No-Gather Export

Layer / File(s)	Summary
Export routing, packing, and configuration modelopt/torch/export/unified_export_hf.py, CHANGELOG.rst, examples/hf_ptq/README.md	The exporter uses the distribution plan and MODELOPT_FSDP2_EXPORT=nogather to select the no-gather path. Packing can retain fused experts. The path writes base and generation configuration, and rejects save_modelopt_state=True for non-offloaded exports. The changelog and README document the option and distribution-plan requirements.
Distributed shard processing and writing modelopt/torch/export/distributed_writer.py	Adds a writer that processes sharded state, synchronizes expert input scales, processes dense and expert weights, assigns tensors to owner ranks, and writes checkpoint shards. Rank 0 adds non-overlapping extra state.
Writer and exporter validation tests/unit/torch/export/test_distributed_writer.py, tests/gpu/torch/export/test_distributed_hf_export.py, tests/unit/torch/export/test_unified_export_hf.py, tests/unit/torch/export/test_fused_experts_keep_fused.py	Adds CPU, multi-rank, and GPU tests for tensor contents, quantization, checkpoint metadata, sidecar files, and shard assembly. Updates source-file-copy coverage and removes earlier helper-level writer checks and the fused-expert event-order test.

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant unified_export_hf
  participant distributed_save_hf_checkpoint
  participant FSDP2 ranks
  unified_export_hf->>distributed_save_hf_checkpoint: packed model and export options
  distributed_save_hf_checkpoint->>FSDP2 ranks: process sharded tensors and assign owners
  FSDP2 ranks->>distributed_save_hf_checkpoint: owned tensors for checkpoint writing

Loading

Merge Risk: Moderate · up to 12267

Ordinary unsharded exports in multi-rank jobs can fail before reaching their existing export path. Preserve that path before merging. CPU-only gloo exports no longer have the previously reported CUDA-device failure.

Pre-merge checks: 5

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The PR changes only modelopt/torch/export/distributed_writer.py and modelopt/torch/export/unified_export_hf.py in the package; their added lines contain no `torch.load(..., weights_only=Fals…
Title check		The title clearly and concisely describes the main change: writing Hugging Face checkpoints for sharded models without gathering the full model.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The PR changes only modelopt/torch/export/distributed_writer.py and modelopt/torch/export/unified_export_hf.py in the package; their added lines contain no torch.load(..., weights_only=False), pickle-enabled NumPy load, hardcoded trust_remote_code=True, eval, exec, or # nosec. The PR does not modify pyproject.toml or requirements files, and it adds no new PIP dependency.

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

[github-actions]

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

[codecov]

Codecov Report

❌ Patch coverage is 94.54545% with 6 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (shengliangx/hfep-12b-writer-pieces@7a249ef). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
modelopt/torch/export/distributed_writer.py	96.15%	3 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	90.62%	3 Missing ⚠️

Additional details and impacted files

@@                          Coverage Diff                          @@
##             shengliangx/hfep-12b-writer-pieces    #2752   +/-   ##
=====================================================================
  Coverage                                      ?   62.30%           
=====================================================================
  Files                                         ?      656           
  Lines                                         ?    73764           
  Branches                                      ?        0           
=====================================================================
  Hits                                          ?    45955           
  Misses                                        ?    27809           
  Partials                                      ?        0           

Flag	Coverage Δ
unit	62.30% <94.54%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 72eab4a

[shengliangxu]

/ok to test a30a335

[shengliangxu]

/ok to test 717b73e

[shengliangxu]

/ok to test ac29323