Repository navigation
Fix the FSDP2 export deadlock at world > 1 - #2731
Merged
Merged
Conversation
A forward leaves the root FSDP module's parameters unsharded, because FSDP2 skips the root's post-forward reshard, and FSDP2's state_dict() pre-hook reshards a module's whole parameter group. In collect_export_tensors only the unit's owner calls state_dict(), so the owner of the non-layer unit resharded the root on its rank alone. The next gather window then all-gathered parameters on that rank that the other ranks still held as plain tensors, a collective they never joined, and the export hung (rank 0 in collect_export_tensors, the rest at the barrier). Both the streaming (owner="share") and rank-0 (owner="rank0") FSDP2 exports were affected. Reshard every FSDP module on every rank before the loop. It is local, and it makes the owner's pre-hook a no-op. The two GPU tests that time out in the nightly 2-GPU run, test_fsdp2_streaming_export_matches_reference[nvfp4|fp8], pass again. Add a two-rank CPU/gloo unit test, which runs in PR CI, so the multi-rank path no longer depends on the nightly to catch a regression. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to The shared export path reshares modules on every rank, and the two-rank streaming test checks completion and output correctness. No concrete merge-blocking risk is established. Pre-merge checks ·
  |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: fixing the FSDP2 export deadlock that occurs with more than one rank. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The authoritative PR diff changes only modelopt/torch/export/unified_export_hf_streaming.pyand a test file. The production addition only iterates FSDP modules and callsmodule.reshard(). It…
- Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
-
📝 Generate docstrings 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2731 +/- ##
==========================================
+ Coverage 71.76% 77.88% +6.11%
==========================================
Files 644 647 +3
Lines 71477 72532 +1055
==========================================
+ Hits 51298 56491 +5193
+ Misses 20179 16041 -4138
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:
|
cjluo-nv
added a commit
that referenced
this pull request
Oct 12, 2026
…port evaluates them (#2710) ### What does this PR do? Type of change: bug fix, documentation **Goal:** evaluate GGML-quantized models (IQ formats, Q8_0) in vLLM, which has no GGML kernels. `--vllm_fakequant_export` already does this. It folds each weight quantizer into the weight and writes a plain BF16 Hugging Face checkpoint with no quantization config, which a stock vLLM image serves as an ordinary model. Its extra `vllm_fq_modelopt_state.pth` is read only by the fake-quant launcher. **The bug:** folding was wrong for GGML weights that GPTQ calibrated, which covers every built-in IQ recipe. - Both fold paths, `mtq.fold_weight` and the vLLM fake-quant export, quantize a float32 copy of the weight. - A BF16 weight holds the pinned payload's decoding rounded to BF16, so that float32 copy never matched the payload GPTQ pinned. - The fold therefore re-encoded the weight and folded in different codes: GPTQ's output quantized a second time, rounding to nearest. - On Qwen3-0.6B with the `iq2_xxs` recipe, this raised vLLM perplexity from 42.2 to 72.0. **Changes:** - **Exact folding:** the GGML pinned-payload lookup also matches a float32 weight against the decoding rounded through BF16 or FP16. Folding a GPTQ'd weight now decodes GPTQ's codes; it doesn't re-encode, and is faster as a result. - **Packed export fix:** GPTQ's pinned payload was also written into the packed export's state dict as a stray `*.weight_quantizer._ggml_pinned_<format>` tensor per layer, which loaders reject. It's now filtered out. - **Docs:** `--vllm_fakequant_export` is documented as the way to evaluate GGML checkpoints in vLLM. An earlier revision of this PR added a separate `upcast_ggml` export option. As review pointed out, folding is the more general mechanism, so that option is gone. The same evaluation now goes through the existing fake-quant export, and the fix above makes folding exact for every caller. ### Usage ```bash python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path <model> --recipe general/ptq/iq2_xxs \ --export_path <dir> --vllm_fakequant_export vllm serve <dir> # stock vLLM image, no quantization argument ``` ### Testing - **New unit test, `test_folding_a_bf16_weight_decodes_the_pinned_payload`:** real GPTQ on a BF16 IQ1_S layer, after which both fold paths reproduce GPTQ's payload exactly. It fails without the fix. - **New unit test, `test_export_drops_the_payload_gptq_pinned`:** runs for all six GGML formats and fails without the skip-key fix. - **Folding check:** folding now matches the decoded packed payload exactly for IQ2_XXS and Q8_0, with and without a pinned payload. Before the fix, 97.9% of pinned IQ2_XXS weights differed. - **End to end (Qwen3-0.6B, `iq2_xxs` recipe, RTX PRO 6000):** `hf_ptq.py --vllm_fakequant_export`, then a stock `vllm/vllm-openai:v0.30.0` container with no modelopt installed. - **Checkpoint:** `config.json` (`Qwen3ForCausalLM`, no `quantization_config`), BF16 safetensors with an index, and the tokenizer files. - **Serving:** `vllm serve` serves it and answers completions. - **Accuracy:** with the same calibration on the same GPU, vLLM perplexity is 42.19 with the fix and 72.0 without. 42.19 is identical to decoding GPTQ's packed payload. - **Memory:** peak extra GPU memory of the pinned-payload lookup for a 17408×5120 weight is 256 MiB for a BF16 weight and 510 MiB for the float32 copy folding passes. Before this PR it was 510 MiB and 1700 MiB, and the float32 copy didn't match at all. Matching row slices (Megatron's gate/up split) now hashes rows a column chunk at a time, with bit-identical keys. For the gate half of a fused 34816×5120 weight, that path drops from 3060 to 597 MiB (BF16) and from 3400 to 1107 MiB (float32 export dtype). - **Unit suites:** `tests/unit/torch/export` plus the GGML quantization tests give 514 passed, 3 failed. The 3 failures are `test_export_diffusers.py` tests that fail the same way on `main` in this environment. - **Pre-commit:** passes. ### 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (GPTQ for the GGML formats, where both bugs live, is unreleased) - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Follows #2665 (GPTQ for GGML formats). Related: #2731 fixes an FSDP2 export deadlock found while testing an earlier revision of this PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * GGML-quantized models can be exported as plain BF16 checkpoints for serving with a stock vLLM image, including weight-only formats such as IQ and Q8_0. * Fake-quant exports preserve decoded GGML weight values across BF16 and FP16 precision conversions, including GPTQ-selected values. * **Documentation** * Updated export and deployment guidance to clarify which vLLM serving paths support these exports. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
meenchen
reviewed
Oct 12, 2026
meenchen
left a comment
meenchen
left a comment
Contributor
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Nudge: the reshard fix looks correct, but the new regression builds its reference while rank 0 still belongs to a two-rank process group.
Needs action:
- Build the unsharded reference on both workers before the nonzero-rank return in
test_fsdp2_parallel_export.py, or destroy the process group first so reference calibration is genuinely single-process.
No action needed:
- The fix normalizes FSDP state on every rank before owner-only
state_dict()calls and covers both ownership modes through the shared collector. - Existing test assertions remain intact. Tests were inspected, not executed.
cjluo-nv
added a commit
that referenced
this pull request
Oct 12, 2026
### What does this PR do? Type of change: documentation (agent instructions) Adds a **Pre-push review pass** to `AGENTS.md`, which `CLAUDE.md` links to. Before asking to push, an agent reviews the whole diff against its base as a skeptical reviewer would and reports what it found. If the pass turns up an opportunity, the agent tells the user and offers to address it before pushing, instead of pushing as-is. The pass checks: - **Performance and memory, especially on GPU,** measured rather than reasoned about: full-size temporaries, repeated expensive work, host syncs on hot paths. - **Distributed and offloaded paths:** every rank issues the same collectives in the same order, and gathered or onloaded weights don't outlive their layer. - **Simplicity:** the smallest diff, reusing an existing mechanism rather than adding a parallel one, no dead code or unrelated refactors, and minimal, backward-compatible API changes. - **Correctness:** validate before mutating, fail loudly rather than silently ignore an option, and cover the layouts the change can meet (dtypes, tied weights, offloaded, FSDP2 and meta tensors, fused or grouped modules). - **Evidence:** each fix has a test that fails without it, hang tests fail fast, paths CI skips get run when touched, and the PR says what was and wasn't run. **Why:** a late review pass on #2710 found several issues that a first implementation passed over: - **Memory:** the GGML pinned-payload lookup peaked at about 3 GB for one weight slice, and chunking its row hashing brought that to 0.6 GB. - **A parallel mechanism:** a separate `upcast_ggml` export option duplicated weight folding. Reusing folding was simpler, once folding was made exact for GPTQ'd weights. - **A hang:** an FSDP2 export deadlock came from a rank-dependent collective (#2731). Writing the pass into the workflow makes it a step every PR gets, not an afterthought. ### Usage N/A (agent workflow). ### Testing `pre-commit run --files AGENTS.md` passes (markdownlint included). Docs only. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: N/A - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (contributor workflow, not user-facing) - Did you get Claude approval on this PR?: N/A 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Added guidance for reviewing changes before requesting a push, including checks for performance, correctness, and simplicity. * The guidance also covers validating changes before they take effect, setting test expectations, and reporting checks that were skipped. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The reference checkpoint was built and exported on rank 0 alone, after rank 1 had returned, while both still belonged to the two-rank process group. Nothing on that path issues a collective today, but one would wait on a partner that has already left. Build and export it on every rank before the sharded export, so any collective it issues has a partner and the deadlock guard covers it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Collaborator
Author
|
Addressed in 143a71c: the reference is now built and exported on both ranks before the sharded export, so any collective its calibration or export issues has a partner, and the 45-second deadlock guard covers it too. The test still fails without the fix (stack in |
meenchen
reviewed
Oct 12, 2026
meenchen
left a comment
meenchen
left a comment
Contributor
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Approve: the prior reference-calibration deadlock concern is resolved, and the shared collector fix has meaningful two-rank regression coverage.
No action needed:
- ✔️ Resolved the critical test concern: both ranks now build and export the reference before the sharded export, within the deadlock guard.
- Existing assertions remain intact; the fixture extension preserves single-process behavior. Tests were inspected, not executed.
Complex PR: 1 existing test file modified or removed. Looping in a human for approval.
sugunav14
approved these changes
Oct 12, 2026
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
3 participants
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.
What does this PR do?
Type of change: bug fix
HF export with FSDP2 deadlocks at 2 or more ranks. It hangs with rank 0 in
collect_export_tensorsand every other rank at the barrier after it. This hits both the streaming export (owner="share") and the gather-to-rank-0 export (owner="rank0"). In practice that'shf_ptq --use_fsdp2with more than one rank, and the nightly 2-GPUgpujob, which has hit its 60-minute limit every night for at least a week:test_fsdp2_streaming_export_matches_reference[nvfp4]and[fp8]each time out at 120s. PR CI runs GPU tests on a 1-GPU runner, where these tests skip, so it never saw this.Root cause:
state_dict()pre-hook on every module that holds FSDP parameters, and the hook reshards that module's whole parameter group.collect_export_tensors, only a unit's owner callsstate_dict(). So the owner of the non-layer unit (embed_tokens,norm,lm_head) reshards the root on its rank alone.A per-rank trace of every collective shows the two ranks agreeing up to that one extra all-gather.
Fix: before the loop, every rank reshards every FSDP module. That's a local operation, so the owner's pre-hook becomes a no-op.
Usage
No API change.
Testing
test_streaming_export_two_cpu_ranksintests/unit/torch/export/test_fsdp2_parallel_export.py. It shards a tiny Llama over two CPU ranks with Gloo, calibrates it (which leaves the root unsharded), runs the streaming export, and compares the merged checkpoint with a single-process export.tests/gpu/torch/export/test_fsdp2_export.pypass in CI's image (nvcr.io/nvidia/pytorch:26.05-py3) on 2× RTX PRO 6000, the nightly runner's GPU. That includestest_fsdp2_streaming_export_matches_reference[nvfp4](4.5s) and[fp8](3.0s), which time out without the fix.test_fsdp2_parallel_export.py(world=1) pass.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Found while adding FSDP2 coverage to an earlier revision of #2710.
🤖 Generated with Claude Code
Summary by CodeRabbit