[yingguo-trt]

What does this PR do?

Type of change: Bug fix

With HuggingFace checkpointing enabled, export_hf_checkpoint(..., save_modelopt_state=False) still writes modelopt_state.pth for Diffusers components. The Diffusers dispatch and component saves omit the flag, so the checkpointing hook uses its own default of True.

Forward the public flag to both quantized and non-quantized model saves. The effective default intentionally changes: omitting the option now follows the public False default. Callers who need the state file can pass True. The same omitted forwarding is present in the 0.47.0 source, so this change includes a changelog entry.

Usage

import modelopt.torch.opt as mto
from modelopt.torch.export import export_hf_checkpoint

mto.enable_huggingface_checkpointing()
export_hf_checkpoint(model, export_dir="export", save_modelopt_state=False)
# Pass save_modelopt_state=True to retain the ModelOpt state file.

Testing

Local offline CPU validation: Windows, Python 3.12.10, PyTorch 2.12.0+cpu, Diffusers 0.40.0, Transformers 5.5.4; OMP_NUM_THREADS=1, matching the unit-test workflow.

pytest tests/unit/torch/export/test_export_diffusers.py::test_export_diffusers_save_modelopt_state -q --durations=4

• 4 passed, each call taking 0.06–0.09s. One parameterized test covers False/True with real FP8 tiny-UNet quantizers enabled or disabled. Disabling the quantizers selects the non-quantized save branch while retaining ModelOpt conversion state. The existing enable_hf_checkpointing fixture enables the save hook.
• Ran the same test against pre-fix source 378888145a in an isolated checkout: both False cases fail because the sidecar is written; both True cases pass.
• Previously ran the original QA CPU reproducer before/after the fix: its direct-save control passes in both versions, while export with False fails before and passes afterward. Config and safetensors files are produced in both versions.
• File-scoped pre-commit checks passed except for the existing Windows Mypy os.sched_getaffinity error in unchanged prepare_megatron_data_blend.py:373, also reproduced on the baseline. The same Mypy version and hook arguments with --platform linux pass for both changed Python files. git diff --check passed.

This validates the state-file saving contract on CPU, not deployment inference or the saved state file's reloadability.

Before your PR is "Ready for review"

Contributor and security guidance reviewed; both commits are signed and signed off.

• Is this change backward compatible?: ❌ The effective default for Diffusers exports with checkpointing enabled changes to omit the sidecar; pass save_modelopt_state=True to retain it. The API signature is unchanged.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A — neither.
• Did you write any new necessary tests?: ✅ One parameterized regression test covers both save paths and detects the pre-fix regression.
• Did you update Changelog?: ✅ Documents the corrected default and how to retain the state file.
• Did you get Claude approval on this PR?: ❌ Addressed the review feedback; awaiting re-review.

Additional Information

Fixes NVBug 6982525. The change suppresses new sidecar writes; it retains the existing handling of files already present in the export directory.

Summary by CodeRabbit

• Bug Fixes
  • Diffusers checkpoint exports now honor the ModelOpt state-saving option: state files are omitted by default and can be retained when explicitly enabled.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: main
• File: .coderabbit.yaml
• Commit: 3788881
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 6c3ce8d7-f125-4922-83d6-1c199cc90993

Commits

Reviewing files that changed from the base of the PR and between 5ff581d and db812a1.

Files selected for processing: 2

• CHANGELOG.rst
• tests/unit/torch/export/test_export_diffusers.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

Diffusers checkpoint export now accepts save_modelopt_state and passes it to component saves. export_hf_checkpoint forwards the option. Tests cover quantized and non-quantized exports with state saving enabled and disabled.

Changes: Diffusers checkpoint state saving

Layer / File(s)	Summary
Configure component state saving modelopt/torch/export/unified_export_hf.py, tests/unit/torch/export/test_export_diffusers.py, CHANGELOG.rst	_export_diffusers_checkpoint accepts save_modelopt_state, defaulting to False. The export path passes the option to quantized and non-quantized component saves. Tests check whether modelopt_state.pth exists for both option values and component cases. The changelog documents the behavior.

Priority: Normal

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

Change: Bug fix

Possibly related PRs

• NVIDIA/Model-Optimizer#2971: Changes Diffusers loading to find the same ModelOpt state file under a component subfolder.

Suggested labels: cherry-pick-0.48.0

Suggested reviewers: kevalmorabia97

Merge Risk: Minimal · up to db812

The export option is correctly propagated and covered for both Diffusers component paths.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

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. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

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: honoring save_modelopt_state during Diffusers export.
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 adds only save_modelopt_state forwarding in modelopt/torch/export/unified_export_hf.py and a test. No added torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True)…

Full details: Security Anti-Patterns

Explanation

PASS. The PR adds only save_modelopt_state forwarding in modelopt/torch/export/unified_export_hf.py and a test. No added torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval/exec, or # nosec bypass appears in the diff. pyproject.toml and requirements.txt are unchanged.

Full details: Docstring Coverage

Explanation

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. (1 skipped: 1 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

🧪 Generate unit tests (beta)

• 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.}

[cjluo-nv]

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

Nudge: the fix looks correct, but it adds no regression test to the repo. The only evidence is a QA test that lives outside the repo and was run by hand.

Needs action:

• Add a regression test to tests/unit/torch/export/test_export_diffusers.py. It should call mto.enable_huggingface_checkpointing(), export a tiny quantized diffusers component with save_modelopt_state=False and assert modelopt_state.pth is absent. A second case with True should assert the file is written. Cover both the quantized and non-quantized save_pretrained calls.
• Confirm the behaviour change is intended: Diffusers exports with HF checkpointing enabled used to write the sidecar by default and now don't. If so, add a CHANGELOG entry so users who rely on it can pass save_modelopt_state=True.

No action needed:

• Passing save_modelopt_state straight to save_pretrained matches the transformers path. The patched _new_save_pretrained pops it, and unpatched diffusers/transformers save_pretrained methods accept extra **kwargs.
• No in-repo diffusers example enables HF checkpointing, so the default change has little practical reach.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.14%. Comparing base (3788881) to head (db812a1).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3012      +/-   ##
==========================================
+ Coverage   77.51%   79.14%   +1.63%     
==========================================
  Files         661      661              
  Lines       74937    74937              
==========================================
+ Hits        58086    59310    +1224     
+ Misses      16851    15627    -1224     

Flag	Coverage Δ
examples-diffusers	19.90% <100.00%> (-0.01%)	⬇️
examples-gpt-oss	13.45% <0.00%> (ø)
examples-hf_ptq	23.34% <0.00%> (-0.04%)	⬇️
examples-llm_distill	13.52% <0.00%> (-0.01%)	⬇️
examples-llm_eval	17.17% <0.00%> (ø)
examples-llm_qat	17.34% <0.00%> (-0.04%)	⬇️
examples-llm_sparsity	16.76% <0.00%> (ø)
examples-megatron_bridge	26.48% <0.00%> (-0.22%)	⬇️
examples-specdec_bench	13.23% <0.00%> (ø)
examples-speculative_decoding	17.42% <0.00%> (-0.06%)	⬇️
examples-torch_onnx	21.22% <0.00%> (-0.02%)	⬇️
examples-torch_trt	15.15% <0.00%> (ø)
examples-vllm_serve	17.38% <0.00%> (ø)
gpu	58.99% <0.00%> (+7.94%)	⬆️
regression	15.04% <0.00%> (-0.01%)	⬇️
unit	61.13% <100.00%> (+<0.01%)	⬆️

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.

[yingguo-trt]

Addressed both items from this review in db812a1.

• Added test_export_diffusers_save_modelopt_state in tests/unit/torch/export/test_export_diffusers.py: one parameterized test runs a real FP8 tiny UNet with False/True through both save branches. The existing enable_hf_checkpointing fixture enables the hook; disabling all quantizers exercises the non-quantized branch while preserving ModelOpt state.
• The effective default change is intended: Diffusers export now honors the public False default. Added a CHANGELOG.rst entry telling users to pass save_modelopt_state=True to retain the state file, and updated the PR description.

Validation: all four cases pass on CPU (0.06–0.09s per call). With the same tests on pre-fix source 378888145a, both False cases fail on the unexpected state file and both True cases pass.

[cjluo-nv]

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

Approving: both concerns from the last review are fixed in the diff, and the flag forwarding is correct and minimal.

No action needed:

• ✔️ Resolved since the last review: the missing regression test (now a parameterized test over both save_pretrained branches, True/False) and the missing CHANGELOG entry for the default change.
• Disabling every quantizer makes has_quantized_modules return false, so the test really runs the non-quantized branch while ModelOpt state is kept. The enable_hf_checkpointing fixture already exists in tests/conftest.py.
• The author ran the test against the pre-fix source: both False cases failed and both True cases passed, so it catches the bug.