[Serialization]: remove explicit weights_only default from safe_load to allow user to bypass if needed

[kevalmorabia97]

Summary

• Remove the kwargs.setdefault("weights_only", True) call from safe_load, deferring to torch's built-in default (which is True for torch>=2.6)
• This allows users to override via the TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 env var when they trust a checkpoint but hit pickle.UnpicklingError
• Add a test that verifies the default fails on unsafe objects and the env var bypass works

Test plan

• python -m pytest tests/unit/torch/utils/test_serialization.py -v

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Serialization utility now respects PyTorch's default behavior and environment-variable configuration instead of forcibly enforcing parameter overrides, providing greater configuration flexibility.

• Tests

  • Added test coverage validating environment-variable override functionality and default behavior in the serialization utility.

[coderabbitai]

📝 Walkthrough

Walkthrough

The safe_load function no longer forcibly sets weights_only=True, instead allowing PyTorch's default behavior or caller overrides to take effect. Tests are added to verify environment variable override functionality.

Changes

Cohort / File(s)	Summary
Core Implementation modelopt/torch/utils/serialization.py	Updated safe_load docstring and removed forced weights_only=True assignment, allowing PyTorch's default behavior to apply.
Test Coverage tests/unit/torch/utils/test_serialization.py	Added test function test_safe_load_env_var_bypasses_weights_only with helper class _UnsafeObj to verify that TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD environment variable successfully bypasses weights-only restrictions.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Security Anti-Patterns	✅ Passed	No critical security anti-patterns were detected in the PR. The changes do not introduce torch.load with weights_only=False, numpy.load with allow_pickle=True, hardcoded trust_remote_code=True, unsafe eval/exec calls, forbidden # nosec comments, or suspicious dependencies.
Title check	✅ Passed	The title accurately describes the main change: removing the explicit weights_only default from safe_load to allow users to bypass it when needed.

_{✏️ 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 safe-load-remove-weights-only-default

---

_{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-04-17 06:25 UTC

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

modelopt/torch/utils/serialization.py (1)

56-65: ⚠️ Potential issue | 🟠 Major

Strengthen safe_load with explicit weights_only=True for defensive clarity.

The codebase enforces torch>=2.8 globally (pyproject.toml line 41), which means torch.load() defaults to weights_only=True when unspecified. However, per SECURITY.md best practices, prefer explicit safe defaults over implicit reliance on PyTorch's version-dependent behavior. Additionally, the current implementation allows callers to override or bypass this safety via environment variables without an explicit guard in the function itself.

🔧 Recommended hardening

 def safe_load(f: str | os.PathLike | BinaryIO | bytes, **kwargs) -> Any:
     """Load a checkpoint securely using ``weights_only=True`` by default.

     NOTE: We dont set default ``weights_only`` (interpret as True for torch>=2.6) so you can override it with
     ``export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1`` if you see ``pickle.UnpicklingError`` and trust the checkpoint.
     """
     if isinstance(f, (bytes, bytearray)):
         f = BytesIO(f)
 
+    # Preserve secure default unless user explicitly overrides or opts out via env var.
+    if "weights_only" not in kwargs and os.getenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD") != "1":
+        kwargs["weights_only"] = True
+
     return torch.load(f, **kwargs)

🤖 Prompt for AI Agents

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

In `@modelopt/torch/utils/serialization.py` around lines 56 - 65, The safe_load
function should explicitly set weights_only=True by default to avoid relying on
torch's version defaults; update safe_load (the function using BytesIO and
calling torch.load) to call kwargs.setdefault('weights_only', True) (so callers
can still explicitly override by passing weights_only) before invoking
torch.load(f, **kwargs), preserving the existing bytes/BytesIO handling.

🤖 Prompt for all review comments with AI agents

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

Outside diff comments:
In `@modelopt/torch/utils/serialization.py`:
- Around line 56-65: The safe_load function should explicitly set
weights_only=True by default to avoid relying on torch's version defaults;
update safe_load (the function using BytesIO and calling torch.load) to call
kwargs.setdefault('weights_only', True) (so callers can still explicitly
override by passing weights_only) before invoking torch.load(f, **kwargs),
preserving the existing bytes/BytesIO handling.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: e38c95c9-3f5b-46d5-bbfa-d28144ac16bf

📥 Commits

Reviewing files that changed from the base of the PR and between 6ded36b and bd889af.

📒 Files selected for processing (2)

• modelopt/torch/utils/serialization.py
• tests/unit/torch/utils/test_serialization.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 76.56%. Comparing base (6ded36b) to head (bd889af).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1279      +/-   ##
==========================================
+ Coverage   75.58%   76.56%   +0.97%     
==========================================
  Files         459      459              
  Lines       48613    48612       -1     
==========================================
+ Hits        36745    37219     +474     
+ Misses      11868    11393     -475     

Flag	Coverage Δ
examples	41.38% <ø> (+11.54%)	⬆️
gpu	59.96% <ø> (-0.54%)	⬇️
unit	52.21% <ø> (-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.

[shengliangxu]

LGTM