[NVBug 6102977] Disable use_cache during export dummy forward for custom configs

[meenchen]

The export step's collect_shared_input_modules() runs a dummy forward pass to trace shared input modules for resmoothing/requantization. This forward call hits the same AttributeError as PR #1324 on configs that don't assign use_cache (e.g., stepfun-ai/Step-3.5-Flash's Step3p5Config):

AttributeError: 'Step3p5Config' object has no attribute 'use_cache'

Wrap the dummy forward with the _disable_use_cache context manager so it sets/restores config.use_cache around the call, mirroring the calibration fix from #1324. Covers all callsites since both unified_export_hf and plugins/vllm_fakequant_hf funnel through collect_shared_input_modules.

What does this PR do?

Type of change: ?

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• 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?: ✅ / ❌ / N/A

Additional Information

Summary by CodeRabbit

• Bug Fixes
  • Improved model export process by optimizing cache handling during profiling operations.

[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: 614b1991-bc52-417f-9687-8a50c74c8c32

📥 Commits

Reviewing files that changed from the base of the PR and between 1c2085f and 52d7a1e.

📒 Files selected for processing (1)

• modelopt/torch/export/unified_export_hf.py

---

📝 Walkthrough

Walkthrough

A single import is added for a cache-disabling utility, and the dummy forward probe in collect_shared_input_modules is wrapped with an additional context manager to disable dataset use_cache during module input collection.

Changes

Cache Disabling During Input Module Collection Probe

Layer / File(s)	Summary
Utility Import modelopt/torch/export/unified_export_hf.py	_disable_use_cache is imported from modelopt.torch.utils.dataset_utils.
Probe Context Wrapping modelopt/torch/export/unified_export_hf.py	Dummy forward probe is nested within _disable_use_cache(model) context manager to prevent cache use during input module collection.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 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 specifically describes the main change: disabling use_cache during export dummy forward for custom configs, which directly addresses the AttributeError fix mentioned in the PR objectives.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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	No security anti-patterns found. PR adds safe context manager usage with no torch.load, numpy.load, trust_remote_code, eval, exec, nosec, or new dependencies.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 weimingc/nvbug_6102977_v2

---

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

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-05-06 23:13 UTC

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 76.82%. Comparing base (097293b) to head (52d7a1e).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1402      +/-   ##
==========================================
+ Coverage   76.73%   76.82%   +0.09%     
==========================================
  Files         476      476              
  Lines       51306    51307       +1     
==========================================
+ Hits        39368    39419      +51     
+ Misses      11938    11888      -50     

Flag	Coverage Δ
examples	41.81% <100.00%> (+2.63%)	⬆️
gpu	59.84% <100.00%> (-0.58%)	⬇️
regression	15.20% <100.00%> (+0.07%)	⬆️
unit	52.54% <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 Sentry.
📢 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 — DM the bot to share feedback.

Tiny, well-targeted bug fix (+5 -1). Reuses the existing _disable_use_cache context manager from dataset_utils to wrap the dummy forward pass in collect_shared_input_modules, mirroring the calibration fix from PR #1324. The context manager correctly handles configs that lack use_cache (sets it to False, restores/deletes on exit). No new code to speak of — just adding one more context manager to an existing with block.

[cjluo-nv]

Bot review — DM the bot to share feedback.

Tiny, correct bug fix (+5 -1). Reuses the existing _disable_use_cache context manager to wrap the dummy forward in collect_shared_input_modules, matching the established pattern from the calibration fix in PR #1324. Handles configs lacking use_cache correctly. No new code, no design concerns.