Use weights_only for torch.load checkpoints - #28097
Merged
Merged
Conversation
Tianlei Wu (tianleiwu)
requested review from
Edward Chen (edgchen1) and
kunal-vaishnavi
April 17, 2026 04:13
kunal-vaishnavi
approved these changes
Apr 17, 2026
Adrian Lizarraga (adrianlizarraga)
added a commit
that referenced
this pull request
May 13, 2026
### Description Follow-up to PR #28097. Applies the same `_torch_load_weights_only()` wrapper to the two remaining `torch.load()` call sites. `torch.load` can deserialize arbitrary Python pickle payloads. Using `weights_only=True` restricts loading to tensor/checkpoint data on supported PyTorch versions and is the safer default. The wrapper gracefully falls back to the default `torch.load` behavior on older PyTorch versions that do not support the `weights_only` parameter. ### Summary of Changes | File | Change | |------|--------| | `onnxruntime/test/testdata/test_data_generation/lr_scheduler/lr_scheduler_test_data_generator.py` | Adds `_torch_load_weights_only()` helper and uses it when loading scheduler/optimizer state dicts. | | `orttraining/orttraining/test/python/orttraining_test_ortmodule_pytorch_ddp.py` | Adds `_torch_load_weights_only()` helper and uses it when loading DDP model checkpoint. | ### Motivation and Context These were the last two `torch.load()` calls in the repository without `weights_only=True`. While both are in test/tooling code with low direct risk, this change ensures consistency with the pattern established in PR #28097 and eliminates all unsafe deserialization call sites. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR updates PyTorch checkpoint loading in the T5 helper and NVIDIA pretraining resume script to prefer
torch.load(..., weights_only=True). Newer PyTorch versions recommend this safer load mode for checkpoints because default pickle-based loading can execute arbitrary code when a.ptfile is malicious.Behavior changes across PyTorch versions:
Summary of Changes
onnxruntime/python/tools/transformers/models/t5/t5_helper.pyweights_only=Truewhen available and uses it forstate_dict_path.orttraining/tools/scripts/nv_run_pretraining.pyckpt_*.pttraining checkpoints.Motivation and Context
torch.loadcan deserialize Python pickle payloads. Usingweights_only=Truenarrows loading to tensor/checkpoint data on supported PyTorch versions and is the safer default for model weights. This reduces risk if an attacker can place or substitute a local.ptfile, or persuades an operator to download and resume from a malicious checkpoint.Testing
python -m py_compile onnxruntime/python/tools/transformers/models/t5/t5_helper.py orttraining/tools/scripts/nv_run_pretraining.py