security(opt): enable weights_only=True by default

[RinZ27]

What does this PR do?

Type of change: Bug fix (Security)

This PR enables weights_only=True by default in load_modelopt_state and restore functions.

The modelopt_state is designed to store metadata using Pydantic's model_dump(), which results in plain Python dictionaries and basic types (lists, strings, numbers). These types are natively supported by PyTorch's weights_only=True mode.

Enabling this by default provides a critical security hardening against potential RCE from malicious state files without breaking existing workflows.

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?: N/A
• Did you update Changelog?: ❌

Additional Information

Related to issue #1055.

Summary by CodeRabbit

• Chores
  • Enhanced checkpoint loading security by restricting model state restoration to weights only, preventing unsafe object deserialization from saved states.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

📝 Walkthrough

Walkthrough

The load_modelopt_state() and restore() functions in the conversion module are updated to load checkpoints with weights_only=True, restricting loaded data to model weights only instead of arbitrary Python objects.

Changes

Cohort / File(s)	Summary
Checkpoint Loading Security modelopt/torch/opt/conversion.py	Updated load_modelopt_state() and restore() functions to set weights_only=True when loading checkpoints, restricting deserialization to model weights only. Comments updated to reflect this parameter change.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title clearly and concisely describes the main security change: enabling weights_only=True by default in model state loading.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Security Anti-Patterns	✅ Passed	PR successfully implements security improvement by enabling weights_only=True by default in load_modelopt_state and restore functions with proper inline security documentation. No new security anti-patterns introduced.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

📝 Coding Plan

• Generate coding plan for human review comments

---

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

[coderabbitai]

🧹 Nitpick comments (2)

modelopt/torch/opt/conversion.py (2)

630-632: Consider adding a security comment for consistency.

The weights_only=True default is correct and aligns with the security guidelines. For consistency with load_modelopt_state (line 526), consider adding a brief inline comment explaining this is a security measure for ModelOpt-generated checkpoints.

📝 Optional: Add security comment for consistency

     # load checkpoint
     kwargs.setdefault("map_location", "cpu")
+    # Security NOTE: weights_only=True is used here on ModelOpt-generated checkpoints
     kwargs.setdefault("weights_only", True)
     objs = torch.load(f, **kwargs)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/opt/conversion.py` around lines 630 - 632, Add a brief inline
security comment next to the kwargs.setdefault("weights_only", True) in
conversion.py (near the torch.load call) explaining that setting
weights_only=True is a security measure for loading ModelOpt-generated
checkpoints, mirroring the comment style used in load_modelopt_state to make
intent consistent and clear.

---

553-553: Documentation example could reinforce secure loading pattern.

The example shows torch.load("model_weights.pt") without weights_only=True. While PyTorch >= 2.6 defaults to True, explicitly showing the secure pattern in documentation helps reinforce best practices for users on older PyTorch versions.

📝 Optional: Update docstring example

-        model.load_state_dict(torch.load("model_weights.pt"), ...)  # Load the model weights
+        model.load_state_dict(torch.load("model_weights.pt", weights_only=True), ...)  # Load the model weights

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/opt/conversion.py` at line 553, Update the docstring/example
to explicitly use the secure loading pattern by passing weights_only=True to
torch.load in the example call (i.e., change the call shown next to
model.load_state_dict to use torch.load("model_weights.pt", weights_only=True));
reference the model.load_state_dict and torch.load usage in conversion.py so
readers see the explicit secure flag.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@modelopt/torch/opt/conversion.py`:
- Around line 630-632: Add a brief inline security comment next to the
kwargs.setdefault("weights_only", True) in conversion.py (near the torch.load
call) explaining that setting weights_only=True is a security measure for
loading ModelOpt-generated checkpoints, mirroring the comment style used in
load_modelopt_state to make intent consistent and clear.
- Line 553: Update the docstring/example to explicitly use the secure loading
pattern by passing weights_only=True to torch.load in the example call (i.e.,
change the call shown next to model.load_state_dict to use
torch.load("model_weights.pt", weights_only=True)); reference the
model.load_state_dict and torch.load usage in conversion.py so readers see the
explicit secure flag.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a5f245a3-81e0-4c09-8930-881be6f870b2

📥 Commits

Reviewing files that changed from the base of the PR and between 00fa5bd and f549221.

📒 Files selected for processing (1)

• modelopt/torch/opt/conversion.py