fix(hf_ptq): use no_grad instead of inference_mode in export_quantized (NVBug 6537702)

[Edwardf0t1]

What does this PR do?

Type of change: Bug fix

Fixes NVBug 6537702 / OMNIML-5658 — multi-node FSDP2 PTQ export fails on all ranks:

hf_ptq.py:910 export_quantized -> export_hf_checkpoint
unified_export_hf.py:1446 _export_transformers_checkpoint -> get_model_state_dict
modelopt/torch/opt/_hooks.py:88 _get_model_state_dict_with_dm_check
torch/distributed/checkpoint/state_dict.py:481 _get_model_state_dict
torch/nn/modules/module.py:2160 _save_to_state_dict
    destination[prefix + name] = param if keep_vars else param.detach()
RuntimeError: Cannot set version_counter for inference tensor

Root cause. export_quantized wrapped its whole body in torch.inference_mode(). On the FSDP2 path (--use_fsdp2), get_model_state_dict(full_state_dict=True) gathers the full params inside that context, so the gathered tensors are inference tensors. Inference tensors have no version counter, so the subsequent state_dict() → param.detach() raises.

Fix. Use torch.no_grad() for the export context. It still disables autograd, but the gathered params stay normal tensors with an intact version counter, so detach() works. FSDP2-only failure — the non-FSDP2 path never hit it because its params already exist outside the context.

The one-line fix is originally by @shengliangx (b0e4328 on shengliangx/distributed-unified); this PR retargets it to the post-rename examples/hf_ptq/ path and adds a changelog entry and a regression guard.

Usage

# 2 nodes x 8 GB200, previously failed at export on every rank
torchrun --nnodes=2 --node_rank=0 --master_addr=$MASTER --master_port=6000 --nproc_per_node=8 \
  hf_ptq.py --model Llama-3.1-8B-Instruct --dataset cnn_dailymail \
  --recipe general/ptq/fp8_default-kv_fp8 --batch_size 8 --calib_size 512 \
  --export_path ./Llama-3.1-8B-Instruct-fp8_default-kv_fp8 --use_fsdp2

Testing

• End-to-end on 2 nodes by @shengliangx on the original branch: dense Qwen3-8B and Qwen3-30B-A3B (MoE) FSDP2 PTQ fp8 checkpoints export successfully.
• Added tests/examples/hf_ptq/test_export_quantized_context.py, a CPU-only guard asserting export_quantized enters torch.no_grad() and not torch.inference_mode(). A functional regression test would need a 2-node FSDP2 job, which CI does not run, so this encodes the invariant instead.
• pre-commit run --files clean on all three changed files.

Reporter (Kenny Kang, GPU SWQA) still needs to confirm on the original 2x8 GB200 Llama-3.1-8B repro.

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

Keyword Committed_ModelOpt_0.46.0 on the bug — should land for 0.46.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Fixed multi-node quantized model exports to prevent runtime errors when gathering and detaching parameters.
  • Improved compatibility with FSDP2 during Hugging Face PTQ exports.

• Tests

  • Added coverage to verify the export process uses the compatible gradient context.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5ddf0149-94ac-4fb6-af9a-be7e1ca1abbd

📥 Commits

Reviewing files that changed from the base of the PR and between c32496e and fb737b5.

📒 Files selected for processing (3)

• CHANGELOG.rst
• examples/hf_ptq/hf_ptq.py
• tests/examples/hf_ptq/test_export_quantized_context.py

🚧 Files skipped from review as they are similar to previous changes (3)

• examples/hf_ptq/hf_ptq.py
• tests/examples/hf_ptq/test_export_quantized_context.py
• CHANGELOG.rst

---

📝 Walkthrough

Walkthrough

export_quantized now uses torch.no_grad() for FSDP2 export. A regression test verifies the context choice, and the changelog documents the fix.

Changes

FSDP2 export

Layer / File(s)	Summary
Export context and regression coverage examples/hf_ptq/hf_ptq.py, tests/examples/hf_ptq/test_export_quantized_context.py, CHANGELOG.rst	The export context uses torch.no_grad() instead of torch.inference_mode(). The test checks this behavior, and the changelog records the fix.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: realasma

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main fix: replacing inference mode with no-grad in export_quantized.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The PR adds only torch.no_grad() and comments/tests; no prohibited deserialization, remote-code, eval/exec, # nosec, or dependency changes are present.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/fsdp2-export-inference-tensor

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 66.96%. Comparing base (151d87d) to head (fb737b5).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2047      +/-   ##
==========================================
+ Coverage   65.45%   66.96%   +1.50%     
==========================================
  Files         521      521              
  Lines       59857    59857              
==========================================
+ Hits        39181    40081     +900     
+ Misses      20676    19776     -900     

Flag	Coverage Δ
examples	43.05% <ø> (+2.09%)	⬆️
unit	55.38% <ø> (-0.02%)	⬇️

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.

[cjluo-nv]

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

One-line fix (torch.inference_mode() → torch.no_grad() in examples/hf_ptq/hf_ptq.py::export_quantized) plus a CHANGELOG entry and a new CPU-only test. The root-cause analysis matches what I see in the tree: the FSDP2 export gathers full params inside the context via get_model_state_dict(full_state_dict=True) (intercepted in modelopt/torch/opt/_hooks.py::_get_model_state_dict_with_dm_check), and inference tensors can't have their version counter set on the subsequent param.detach(). no_grad still disables autograd and the export path does no training-style work, so behavior for the non-FSDP2 paths should be unchanged. Diff is tiny, backward compatible, and the new test file's header matches the canonical LICENSE_HEADER (no licensing concern). No prompt-injection content in the PR metadata.

Flagging for owner sign-off on the points below rather than approving, since the actual failure mode can't be exercised in CI and the added test only checks source text.

[sugunav14]

nit: could we trim down the comment?

Suggested change

	# Use no_grad rather than inference_mode: on the FSDP2 path the export gathers the
	# full params inside this context, and tensors created under inference_mode are
	# inference tensors whose version counter cannot be set, so the subsequent
	# state_dict() -> param.detach() fails with
	# "RuntimeError: Cannot set version_counter for inference tensor".
	# FSDP2 export path gathers full params in this context and inference tensors break state_dict()'s
	# detach here.

[Edwardf0t1]

Done — trimmed to two lines in c32496e. Also rebased onto current main to resolve a CHANGELOG conflict from #2052 (the 0.46 section refactor); hf_ptq.py auto-merged cleanly.

[sugunav14]

LGTM, left a minor comment

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-08-04 20:57 UTC