Distributed no-gather HF checkpoint export for FSDP2 by shengliangxu · Pull Request #2228 · NVIDIA/Model-Optimizer

shengliangxu · 2026-08-21T20:32:03Z

What does this PR do?

Type of change: New feature + bug fix

Feature — export an FSDP2 model without gathering it. Exporting an FSDP2-sharded model with export_hf_checkpoint gathered the full unsharded state dict to host RAM on rank 0 and wrote it from there. That puts the whole model on one host, which is both the wall-clock bottleneck and a hard ceiling: large MoE checkpoints OOM at the gather no matter how many nodes are added.

This routes an FSDP2 model under torch.distributed through a per-rank writer instead. Every rank writes only its own weight shards via torch DCP (distributed_save_hf_checkpoint), the output files are consolidated in parallel across ranks, and rank 0 writes config.json, the quant config, and the weight index. No rank ever holds the full model.

The MoE path matters as much as the writer. The in-model _export_fused_experts splits the fused weight into full per-expert tensors and leaves them all in the state dict, so the write point holds every expert on every rank. _export_fused_experts_keep_fused folds the per-expert quantizers into the fused weight in place instead, leaving one fused parameter per projection, and the per-expert split then happens inside the writer on whatever each rank owns. That is what keeps the write-point state dict sharded — measured at 235B as 157 GB -> 22 GB per rank. The single-process and gather paths are unchanged.

Bug fix — the gather path writes the model's unquantized weights into the quantized checkpoint. The exporter splits each fused 3-D expert weight into per-expert quantized weights, but the original mlp.experts.gate_up_proj / mlp.experts.down_proj parameters stay on the module and are picked up by the state-dict gather, so they get written as well: BF16, pre-quantization, carrying no weight_scale/input_scale, into a checkpoint whose hf_quant_config.json declares FP8. They sit under exactly the names transformers binds directly, since its in-memory Qwen3MoeExperts parameters are the fused 3-D tensors.

On Qwen3-30B-A3B that is 96 stale tensors, and experts are ~29B of the model's 30.5B parameters — so the FP8 export ships a full unquantized copy of essentially the whole model: 58.0 GB of an 89.2 GB file that should be 31.2 GB (56307 tensors vs 56211). The no-gather path cannot produce this: _export_fused_experts_keep_fused folds the quantizers into the fused parameter in place, so the parameter holds quantized data rather than the original, and the writer then splits it into the per-expert keys — no fused key survives.

Commits, in order: the writer, routing FSDP2 export through it, the shard-local fused-experts fold, an in-place re-materialization of the write-point state dict, the weight index, and the test.

Usage

# No API change. An FSDP2 model under torch.distributed takes the new path automatically;
# a single-process model is unaffected.
from modelopt.torch.export import export_hf_checkpoint

export_hf_checkpoint(model, export_dir=export_dir)

Testing

New GPU test — tests/gpu/torch/export/test_distributed_hf_export.py. PTQs a tiny Qwen3 (dense and MoE) under FSDP2 across the whole world with the default FP8 and NVFP4 configs, then asserts against the safetensors headers on disk: every source parameter survives and nothing extra is left behind (checked both directions); each quantized weight carries the scales its format needs; every tensor is at its full unsharded shape, accounting for NVFP4's uint8 packing; and the exporter's own files plus the weight index are present and consistent with the tensors on disk, while the sidecars it must leave alone survived. It asserts is_fsdp2_model() inside the workers so the test cannot pass by silently falling back to the gather path. 4 passed in ~32s on 4x GB200.

Real model — Qwen3-30B-A3B, FP8 weights + FP8 KV, 2 nodes x 4 GPUs (world 8). Both sides ran the identical command; only $MODELOPT_SRC differs (base commit vs this branch), plus the output directory:

# $MODELOPT_SRC  - ModelOpt checkout under test
# $EXPORT_DIR    - output directory on shared storage
# $MASTER_ADDR / $MASTER_PORT / $NODE_RANK - supplied by the job launcher, one task per node
export PYTHONPATH="$MODELOPT_SRC:$MODELOPT_SRC/examples/hf_ptq:$PYTHONPATH"

torchrun \
    --nnodes 2 --nproc_per_node 4 --node_rank "$NODE_RANK" \
    --master_addr "$MASTER_ADDR" --master_port "$MASTER_PORT" \
    "$MODELOPT_SRC/examples/hf_ptq/hf_ptq.py" \
        --pyt_ckpt_path Qwen/Qwen3-30B-A3B \
        --recipe general/ptq/fp8_default-kv_fp8 \
        --batch_size 2 \
        --calib_size 64 \
        --calib_seq 512 \
        --dataset cnn_dailymail \
        --export_path "$EXPORT_DIR" \
        --use_fsdp2 \
        --skip_generate \
        --trust_remote_code

phase	gather (base)	this PR
load	53s	51s
quantize + calibrate	68s	67s
export	301s	81s

Load and calibrate are unchanged, as expected — only the export path differs. Of the base's 301s, ~281s is the rank-0 gather and only ~15s is writing, so the gather is 94% of the cost and this removes it.

The same two runs are the measurement behind the bug fix above: 89.2 GB / 56307 tensors on the base commit against 31.2 GB / 56211 here, the difference being the 96 unquantized fused expert tensors. Both MoE cases of the new test fail on the base commit for that reason; both dense cases pass, since dense models have no fused experts to leave behind.

Larger-MoE validation from the fused-fold commit (NVFP4, calib 1): Qwen3-30B-A3B export 262s -> 62s with a byte-identical checkpoint, and Qwen3-235B-A22B exports in 343s (write-point base 157 GB -> 22 GB) where it previously OOM'd on both 16 and 32 GPUs.

Before your PR is "Ready for review"

Is this change backward compatible?: ✅
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?: ✅
Did you get Claude approval on this PR?: ❌

Additional Information

Known gaps in the no-gather path. It auto-triggers for every FSDP2 distributed export, so each of these is a silent behaviour change for callers that hit it, not an opt-in. Three were raised in review and are now fixed; two remain open:

~~extra_state_dict is ignored.~~ Fixed — now raises NotImplementedError rather than dropping the tensors silently.
~~save_modelopt_state=True is ignored.~~ Fixed — now raises NotImplementedError, mirroring how the streaming-offload path already declines it.
~~Tied-weight dedup is not applied.~~ Fixed — the writer now builds a TiedWeightMap from the model and passes it to both postprocess_state_dict calls, restoring the name-based drop from Name-based tied-weight deduplication during HF checkpoint export (supersedes #2092) #2194 for FSDP2 models. Covered by a tied-embedding test that was verified to fail without the change.
Non-MoE conversion_mapping reversals are still skipped. The fused -> per-expert un-fusing happens in the writer, but the quant-aware reverse of any other transformers>=5 key rename runs only on the gather path, so a model relying on those may export in-memory names instead of hub-compatible ones. Still open — likely wants a gate that refuses the fast path when a non-MoE mapping applies.
The fold does not actually quantize only local experts. _export_fused_experts_keep_fused has a Shard(0) DTensor branch for that, but under this path it does not engage: fsdp2_aware_weight_update(..., reshard=False) unshards the enclosing decoder layer before the handler runs, so the fold sees a full tensor and quantizes all of that layer's experts on every rank (measured: DTensor=False, 8 of 8 experts at world 4). Peak memory is still bounded to one layer, because _process_quantized_modules reshards the previous layer before moving on, and the write-point saving is unaffected — but the DTensor branch is effectively dead code here and should be made reachable or removed.

Also: the weight index is built by the writer rather than delegated to torch's consolidation helper. consolidate_safetensors_files_on_every_rank does not emit model.safetensors.index.json on torch 2.11.0a0, and without it a multi-shard checkpoint is unloadable with no error anywhere — worth confirming the behaviour on whichever torch CI pins.

Commits carry DCO Signed-off-by but are not GPG-signed.

Add the shard-local, no-rank-0-gather HF safetensors export, decoupled from any EP/session substrate -- for a FULLY FSDP2-sharded model (no expert parallelism): - _export_common.py (new): _size_to_bytes leaf helper. - moe_utils.py: fused-experts HF-format helpers (_fused_experts_prefixes, split_fused_experts_state_dict, _dtensor_dim0_offset, _split_local_fused_module, _FUSED_PROJ) -- split a rank's LOCAL FSDP2 shard of a fused 3-D expert weight into per-expert HF keys, using the DTensor's own dim-0 offset as the global expert index (ep_rank=0, so every rank writes its own experts). - distribute.py: distributed_save_hf_checkpoint (+ _bin_pack_fqn_to_index, _finfo_accepts_int_dtypes) -- DCP HuggingFaceStorageWriter(save_distributed=True) per-rank write + parallel consolidation, no full-model host-RAM gather. The caller (hf_ptq) exports under torch.inference_mode(), so the writer runs under `inference_mode(False) + no_grad()` and re-materializes each state-dict value into a fresh normal tensor up front (empty_like+copy_; DTensor via to_local->from_local to keep mesh/placements) -- inference tensors have no version counter, which the fused-experts split's slicing and DCP's SavePlanner both require. Validated on hecate, 2 nodes x 4 GPU: Qwen3-0.6B FP8 (1 shard) and Qwen3-30B-A3B FP8 (3 shards, 128 experts sharded 16/rank, correct per-expert global indices 0..127) -- both written entirely no-gather. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> (cherry picked from commit 6737e0d172f67d81ba1499aafa4e6782dc5b2d46)

Replace the FSDP2 gather-to-rank-0 export with the per-rank DCP write for a distributed run. _export_transformers_checkpoint gains defer_distributed_fsdp2_write (default False; only export_hf_checkpoint sets it): when the model is FSDP2 under a live process group it does the quant processing but returns a None state dict instead of gathering the full model to CPU on rank 0. export_hf_checkpoint then has all ranks call distributed_save_hf_checkpoint (per-rank shard write, no host-RAM gather); rank 0 writes config.json + generation_config (_write_base_config) and folds the deployment quant/sparse config into config.json -- mirroring the gather-path tail. hf_spec_export's callers keep the gather behavior. Known open items (flagged in-code): the distributed path does not run revert_weight_conversion_quant_aware (non-expert transformers>=5 key renames), and extra_state_dict is not merged into the distributed write. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> (cherry picked from commit 29383f8b7613cca66e9677ef52e512d8d952af38)

The no-gather distributed export ran the in-model _export_fused_experts, which under FSDP2 all-gathers the Shard(0) fused expert weight (via the export handler's unshard) and then materializes full per-expert weights for ALL experts on EVERY rank -- de-sharding the whole model. Fine at 30B (fits) but OOMs large MoE: at 235B the write-point state dict was ~157GB/rank (~113GB of replicated per-expert weights), and the footprint did not shrink with more nodes (32-GPU OOM'd at the same point as 16-GPU). Add _export_fused_experts_keep_fused: fold the per-expert quantizers into the FUSED 3-D weight in place. On a Shard(0) DTensor it quantizes ONLY this rank's local experts (to_local() + the DTensor dim-0 global offset for the matching per-expert weight quantizer) and rewraps the result Shard(0), so the model stays sharded and DCP writes per-rank shards; the per-expert split then happens on the sharded fused weight (split_fused_experts_state_dict). On a plain tensor (single process) it quantizes all experts -- byte-identical to _export_fused_experts. Dispatched via a keep_fused_experts flag (ExportContext -> the fused-experts export handler), set only when deferring to the no-gather writer; the single-process / gather paths keep the split-in-place behavior unchanged. Validated on hecate (NVFP4, calib 1): Qwen3-30B-A3B export drops 262s -> 62s with a byte-identical checkpoint (74401 keys, 128 experts, 17GB); Qwen3-235B-A22B now exports in 343s (write-point base 157GB -> 22GB) where it previously OOM'd on both 16 and 32 GPUs. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> (cherry picked from commit 234024a20d3eba09be4e49d5c68eb878d12a83ea)

distributed_save_hf_checkpoint rebuilt the sharded state dict as a dict comprehension `{k: _to_normal(v) for ...}`, which keeps the ENTIRE original dict alive while building a full second copy -- a transient 2x of the write-point state dict. Overwrite each key in place instead, so each old value is freed as it is replaced and peak stays ~1x. Same for the per-expert local_sd detach pass. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> (cherry picked from commit 3147fe86523cbdd98b107c37bdb54225e7bfa0af)

…2 writer The no-gather distributed export left the checkpoint without a weight index. distribute.py delegates the index to torch's consolidate_safetensors_files_on_every_rank ("rank 0 writes only the index"), but that helper does not emit model.safetensors.index.json -- on torch 2.11.0a0+eb65b36914.nv26.02 it writes only the merged shards. Without the index transformers / vLLM cannot map keys to shards, so every multi-shard export was silently unloadable: a Qwen3-30B-A3B FP8 run produced 4 correct safetensors, no index, and no error anywhere in the log. Build the index on rank 0 after consolidation. No extra collective is needed -- fqn_to_index_mapping and all_sizes are already unioned across ranks by the all_gather_object above and are identical on every rank -- so rank 0 maps each FQN to the file the writer placed it in and sums the global byte sizes. This also gives n_files (computed and unused until now) its purpose. The shard names must match what HuggingFaceStorageWriter derived from the same mapping: model-<i>-of-<n>.safetensors, 5-digit zero-padded. The ep-dp branch's _write_even_shards_no_consolidation does this with two all_gather_object calls because there each rank knows only its own shards; here the mapping is already global, so the gather is unnecessary. Validated on lyris (Qwen3-30B-A3B, FP8 weights + FP8 KV, 2 nodes x 4 GPU, pure FSDP2, job 2755390): the index covers all 56211 keys with 0 missing, 0 extra and 0 file mismatches when checked against the actual safetensors headers; total_size 31167395712 sits one header-size below the on-disk total, matching the gather path's convention. The weight shards are byte-identical in size to the pre-fix run, so the change adds the index without perturbing the write. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…eckpoint The distributed writer rebuilds a checkpoint from per-rank DCP shards, so the ways it can go wrong are structural, not numerical: a key silently dropped, a rank-local shard written where the full tensor belongs, a sidecar or the weight index never emitted. All of those produce a checkpoint that looks fine -- no error, plausible file sizes -- but is wrong or unloadable. Nothing covered that. Add a GPU test that PTQs a tiny Qwen3 (dense and MoE) under FSDP2 across the whole world with the default FP8 and NVFP4 configs, then asserts against the safetensors headers on disk: 1. every source parameter survives the round trip, and nothing extra is left behind -- checked both directions, so a fused expert weight left next to the per-expert weights it was split into is a failure, not a silent 2.9x checkpoint; 2. each quantized weight carries the scales its format needs (weight_scale + input_scale, plus weight_scale_2 for NVFP4); 3. every tensor is at its full unsharded shape, with NVFP4's uint8 packing accounted for; 4. the files the exporter writes are present (config, generation config, quant config, and the weight index whenever the layout is sharded, cross-checked against the tensors actually on disk), and the sidecars it must leave alone -- tokenizer/vocab/chat template -- survived. The test asserts is_fsdp2_model() inside the workers: the no-gather path is selected by exactly that predicate, and without the check the export would quietly fall back to the rank-0 gather and the test would pass while never touching the code under test. Model dimensions are set explicitly rather than taking the tiny-Qwen3 defaults. Those give head_dim = hidden_size / num_heads = 32/16 = 2, and a 2-element q_norm split over a 4-rank world leaves ranks holding empty chunks, which the DCP planner rejects outright ("invalid fill tensor-volume") and which kills the whole worker pool. Every sharded axis is now comfortably divisible and last dims are multiples of 16 so NVFP4 blocks are well defined. Verified on lyris (4x GB200): 4 passed in ~32s on this branch. Against main (a2fbac7) the two MoE cases fail on check 1 with the leftover fused experts -- the toy-model form of the 96 stale BF16 tensors that made a real Qwen3-30B-A3B FP8 export 89.2 GB instead of 31.2 GB -- while both dense cases pass, so the test discriminates rather than failing blanket. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

New feature entry for the 0.47 section: export_hf_checkpoint now routes an FSDP2 model under torch.distributed through the per-rank DCP writer instead of the rank-0 host-RAM gather. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

copy-pr-bot · 2026-08-21T20:32:07Z

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-08-21T20:32:29Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 638b721d-250a-4219-9dac-31e09df061b8

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

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

github-actions · 2026-08-21T20:36:21Z

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2228/
Built to branch `gh-pages` at 2026-08-21 22:25 UTC. Preview will be ready when the GitHub Pages deployment is complete.

…xport Bug Fixes entry for 0.47: the FSDP2 HF export split each fused expert weight into per-expert quantized weights but left the original unquantized fused tensors in the checkpoint, so every expert was written twice. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

codecov · 2026-08-21T20:50:19Z

Codecov Report

❌ Patch coverage is 5.73770% with 345 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.55%. Comparing base (a2fbac7) to head (51a8fde).

Files with missing lines	Patch %	Lines
modelopt/torch/export/moe_utils.py	3.66%	184 Missing ⚠️
modelopt/torch/export/distribute.py	6.77%	110 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	4.76%	40 Missing ⚠️
modelopt/torch/export/_export_common.py	20.00%	8 Missing ⚠️
modelopt/torch/export/hf_export_handlers.py	25.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2228      +/-   ##
==========================================
- Coverage   78.99%   78.55%   -0.45%     
==========================================
  Files         522      523       +1     
  Lines       60599    60960     +361     
==========================================
+ Hits        47872    47888      +16     
- Misses      12727    13072     +345

Flag	Coverage Δ
unit	`55.31% <5.73%> (-0.30%)`	⬇️

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.

… duplicates Reword the 0.47 bug-fix entry: the leftover mlp.experts.gate_up_proj / down_proj tensors are the model's original BF16 pre-quantization parameters, carrying no scales, so a quantized export shipped a full unquantized copy of every expert rather than a redundant copy of the quantized data. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The no-gather path auto-engages for every FSDP2 export under torch.distributed, so anything the gather path does that it does not is a silent behaviour change for callers who never opted in. Review of #2228 found three; this closes them. extra_state_dict and save_modelopt_state were accepted and ignored. The gather path merges the former into the exported state dict and forwards the latter to save_pretrained; the distributed path writes the checkpoint itself and returns before both, so a caller passing either got a checkpoint quietly missing what they asked for. Reject them with NotImplementedError instead, mirroring how the streaming-offload path already declines save_modelopt_state. The check sits outside the export try/except so the failure surfaces as itself rather than behind the generic "Cannot export model to the model_config" warning, and is gated on `not _offloaded` so an offloaded model still reaches the streaming path. Tied-weight dedup was not applied. The gather path passes a TiedWeightMap into postprocess_state_dict; the writer called it positionally, leaving tied_map=None, so the name-based drop added in #2194 silently did not run for FSDP2 models. fully_shard splits a shared nn.Parameter into distinct per-module shards, so a declared tie reaches the writer as two independent DTensors and both were written. Build the map from the model inside the writer and pass it to both postprocess calls. The map is name-derived and identical on every rank, so the alias is dropped consistently, and postprocess already skips a group whose canonical is absent from this rank's slice rather than orphaning it. Tests: parametrized rejection cases for both options, and a tied-embedding export asserting the alias is gone and the canonical kept. Verified the tied test discriminates -- with tied_map removed from the two call sites it fails with "tied alias 'lm_head.weight' was written alongside its canonical", and passes with it. 7 passed on 4x GB200. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu added 7 commits August 21, 2026 09:50

shengliangxu added 2 commits August 21, 2026 14:21

Navigation Menu

Uh oh!

Distributed no-gather HF checkpoint export for FSDP2 - #2228

Distributed no-gather HF checkpoint export for FSDP2#2228
shengliangxu wants to merge 10 commits into
mainfrom
shengliangx/distributed-export

shengliangxu commented Aug 21, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 21, 2026

Uh oh!

coderabbitai Bot commented Aug 21, 2026 •

edited

Loading

Review skipped

Uh oh!

github-actions Bot commented Aug 21, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-08-21 22:25 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

codecov Bot commented Aug 21, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

shengliangxu commented Aug 21, 2026 • edited 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

Uh oh!

copy-pr-bot Bot commented Aug 21, 2026

Uh oh!

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

Uh oh!

Review skipped

Uh oh!

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

Uh oh!

Built to branch gh-pages at 2026-08-21 22:25 UTC. Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

shengliangxu commented Aug 21, 2026 •

edited

Loading

coderabbitai Bot commented Aug 21, 2026 •

edited

Loading

github-actions Bot commented Aug 21, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-08-21 22:25 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

codecov Bot commented Aug 21, 2026 •

edited

Loading