From bd889afd4fe6c08f1c8054145cb1b4a7f0920524 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 16 Apr 2026 11:52:11 -0700 Subject: [PATCH] feat(serialization): remove explicit weights_only default from safe_load Allow torch>=2.6's built-in default (weights_only=True) to take effect naturally, so users can override via TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 when they trust a checkpoint but hit pickle.UnpicklingError. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/utils/serialization.py | 6 +++-- tests/unit/torch/utils/test_serialization.py | 24 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/serialization.py b/modelopt/torch/utils/serialization.py index da16f7514cc..dc880b86b80 100644 --- a/modelopt/torch/utils/serialization.py +++ b/modelopt/torch/utils/serialization.py @@ -54,9 +54,11 @@ def safe_save(obj: Any, f: str | os.PathLike | BinaryIO, **kwargs) -> None: def safe_load(f: str | os.PathLike | BinaryIO | bytes, **kwargs) -> Any: - """Load a checkpoint securely using weights_only=True by default.""" - kwargs.setdefault("weights_only", True) + """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) diff --git a/tests/unit/torch/utils/test_serialization.py b/tests/unit/torch/utils/test_serialization.py index 32851d3a09e..cb3233739cd 100644 --- a/tests/unit/torch/utils/test_serialization.py +++ b/tests/unit/torch/utils/test_serialization.py @@ -16,7 +16,9 @@ """Tests for Modelopt's serialization utilities.""" from io import BytesIO +from pickle import UnpicklingError +import pytest import torch from modelopt.torch.opt.config import ModeloptBaseConfig @@ -70,3 +72,25 @@ def test_safe_load_with_path(tmp_path): loaded_state = safe_load(file_path) assert loaded_state["data"] == 42 + + +class _UnsafeObj: + """Not registered in torch safe globals — unpickling fails with weights_only=True.""" + + def __init__(self, v): + self.v = v + + +def test_safe_load_env_var_bypasses_weights_only(tmp_path, monkeypatch): + """Verify TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 allows safe_load to load objects unsafe for weights_only.""" + file_path = tmp_path / "unsafe.pt" + torch.save({"obj": _UnsafeObj(42)}, file_path) + + # Always fails when weights_only is not set (default=True) + with pytest.raises(UnpicklingError): + safe_load(file_path) + + # With the env var, safe_load (no explicit weights_only) defers to torch's default=False + monkeypatch.setenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1") + loaded = safe_load(file_path) + assert loaded["obj"].v == 42