From 406a89bf7d1a08f7475d4ebce24297f345bc39c5 Mon Sep 17 00:00:00 2001 From: adrianlizarraga Date: Fri, 8 May 2026 14:39:14 -0700 Subject: [PATCH 1/2] Use weights_only=True for remaining torch.load() calls Add _torch_load_weights_only() wrapper to lr_scheduler_test_data_generator.py and orttraining_test_ortmodule_pytorch_ddp.py, matching the pattern from PR #28097. This mitigates arbitrary code execution risk from malicious PyTorch checkpoints. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../lr_scheduler_test_data_generator.py | 18 +++++++++++++++++- .../orttraining_test_ortmodule_pytorch_ddp.py | 17 ++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py b/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py index 9d59af54206db..c78ad9378845a 100644 --- a/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py +++ b/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py @@ -4,9 +4,25 @@ """This file is used to generate test data for LR scheduler optimizer tests in orttraining/orttraining/test/training_api/core/training_api_tests.cc.""" +import logging + import torch from torch.optim.lr_scheduler import LambdaLR +logger = logging.getLogger(__name__) + + +def _torch_load_weights_only(path: str, **kwargs): + try: + return torch.load(path, weights_only=True, **kwargs) + except TypeError: + logger.warning( + "Current PyTorch version does not support torch.load(..., weights_only=True); " + "falling back to default torch.load behavior for %s.", + path, + ) + return torch.load(path, **kwargs) + class SingleParameterModule(torch.nn.Module): """A dummy module containing only one trainable parameter.""" @@ -90,7 +106,7 @@ def main(): json.dump(data, f, ensure_ascii=False, indent=4) data = [] - state_dict = torch.load(fp.name) + state_dict = _torch_load_weights_only(fp.name) new_adamw_optimizer = torch.optim.AdamW(pt_model.parameters(), lr=1e-3) new_adamw_optimizer.load_state_dict(state_dict["optimizer"]) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py index bb0fedb4938a1..0b8416fd58555 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py @@ -1,6 +1,7 @@ # This test script is a modified version of Pytorch's tutorial. # For details, see https://pytorch.org/tutorials/intermediate/ddp_tutorial.html. import argparse +import logging import os import sys # noqa: F401 import tempfile @@ -15,6 +16,20 @@ import onnxruntime # noqa: F401 from onnxruntime.training.ortmodule import ORTModule +logger = logging.getLogger(__name__) + + +def _torch_load_weights_only(path: str, **kwargs): + try: + return torch.load(path, weights_only=True, **kwargs) + except TypeError: + logger.warning( + "Current PyTorch version does not support torch.load(..., weights_only=True); " + "falling back to default torch.load behavior for %s.", + path, + ) + return torch.load(path, **kwargs) + def setup(rank, world_size): os.environ["MASTER_ADDR"] = "localhost" @@ -113,7 +128,7 @@ def demo_checkpoint(rank, world_size, use_ort_module): dist.barrier() # configure map_location properly map_location = {"cuda:0": f"cuda:{rank}"} - ddp_model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=map_location)) + ddp_model.load_state_dict(_torch_load_weights_only(CHECKPOINT_PATH, map_location=map_location)) optimizer.zero_grad() outputs = ddp_model(torch.randn(20, 10)) From 8da88d8752ad9af7f2c43b0519960bad8a5c3be8 Mon Sep 17 00:00:00 2001 From: adrianlizarraga Date: Fri, 8 May 2026 14:59:23 -0700 Subject: [PATCH 2/2] Address review: use inspect.signature(torch.load) --- .../lr_scheduler_test_data_generator.py | 19 +++++++++++-------- .../orttraining_test_ortmodule_pytorch_ddp.py | 19 +++++++++++-------- 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py b/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py index c78ad9378845a..9f7343133ac26 100644 --- a/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py +++ b/onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py @@ -4,6 +4,7 @@ """This file is used to generate test data for LR scheduler optimizer tests in orttraining/orttraining/test/training_api/core/training_api_tests.cc.""" +import inspect import logging import torch @@ -11,17 +12,19 @@ logger = logging.getLogger(__name__) +_TORCH_LOAD_HAS_WEIGHTS_ONLY = "weights_only" in inspect.signature(torch.load).parameters + def _torch_load_weights_only(path: str, **kwargs): - try: + if _TORCH_LOAD_HAS_WEIGHTS_ONLY: return torch.load(path, weights_only=True, **kwargs) - except TypeError: - logger.warning( - "Current PyTorch version does not support torch.load(..., weights_only=True); " - "falling back to default torch.load behavior for %s.", - path, - ) - return torch.load(path, **kwargs) + + logger.warning( + "Current PyTorch version does not support torch.load(..., weights_only=True); " + "falling back to default torch.load behavior for %s.", + path, + ) + return torch.load(path, **kwargs) class SingleParameterModule(torch.nn.Module): diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py index 0b8416fd58555..81e5662a11bc7 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py @@ -1,6 +1,7 @@ # This test script is a modified version of Pytorch's tutorial. # For details, see https://pytorch.org/tutorials/intermediate/ddp_tutorial.html. import argparse +import inspect import logging import os import sys # noqa: F401 @@ -18,17 +19,19 @@ logger = logging.getLogger(__name__) +_TORCH_LOAD_HAS_WEIGHTS_ONLY = "weights_only" in inspect.signature(torch.load).parameters + def _torch_load_weights_only(path: str, **kwargs): - try: + if _TORCH_LOAD_HAS_WEIGHTS_ONLY: return torch.load(path, weights_only=True, **kwargs) - except TypeError: - logger.warning( - "Current PyTorch version does not support torch.load(..., weights_only=True); " - "falling back to default torch.load behavior for %s.", - path, - ) - return torch.load(path, **kwargs) + + logger.warning( + "Current PyTorch version does not support torch.load(..., weights_only=True); " + "falling back to default torch.load behavior for %s.", + path, + ) + return torch.load(path, **kwargs) def setup(rank, world_size):