Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update pytorch weights_only parameter to False #114

Open
wants to merge 2 commits into
base: master
Choose a base branch
from

Conversation

adarshan-intel
Copy link
Contributor

@adarshan-intel adarshan-intel commented Feb 4, 2025


This change is Reviewable

Copy link
Contributor

@kailun-qin kailun-qin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 1 of 1 files at r1, all commit messages.
Reviewable status: all files reviewed, 3 unresolved discussions, not enough approvals from maintainers (3 more required), not enough approvals from different teams (2 more required, approved so far: ) (waiting on @adarshan-intel)


-- commits line 2 at r1:
-> Set weights_only=False when loading models from a trusted source

Code quote:

Update pytorch weights_only parameter to False

-- commits line 5 at r1:
Just give some context that for PyTorch 2.6 and later, direct unpickling with weights_only=True (which was made the default) can fail if the model file contains more than just the weights and includes classes or functions that are not allowlisted.

Code quote:

  - Updated pytorchexample.py to set the weights_only argument to False when calling torch.load.
  - This change resolves the UnpicklingError encountered when loading the pre-trained model.

pytorch/pytorchexample.py line 8 at r1 (raw file):

# Load the model from a file
alexnet = torch.load("alexnet-pretrained.pt", weights_only=False)

Since we are explicitly setting weights_only=False here, can we at least add a note indicating that the file is from a trusted source

Alternatively, could we use weights_only=False for older versions, and for newer versions, we use weights_only=True + safe_globals (to allow only the necessary classes/functions)?

Code quote:

alexnet = torch.load("alexnet-pretrained.pt", weights_only=False)

Copy link
Contributor Author

@adarshan-intel adarshan-intel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewable status: all files reviewed, 3 unresolved discussions, not enough approvals from maintainers (3 more required), not enough approvals from different teams (2 more required, approved so far: ) (waiting on @kailun-qin)


-- commits line 2 at r1:

Previously, kailun-qin (Kailun Qin) wrote…

-> Set weights_only=False when loading models from a trusted source

Done.


-- commits line 5 at r1:

Previously, kailun-qin (Kailun Qin) wrote…

Just give some context that for PyTorch 2.6 and later, direct unpickling with weights_only=True (which was made the default) can fail if the model file contains more than just the weights and includes classes or functions that are not allowlisted.

Done.


pytorch/pytorchexample.py line 8 at r1 (raw file):

Previously, kailun-qin (Kailun Qin) wrote…

Since we are explicitly setting weights_only=False here, can we at least add a note indicating that the file is from a trusted source

Alternatively, could we use weights_only=False for older versions, and for newer versions, we use weights_only=True + safe_globals (to allow only the necessary classes/functions)?

Done.

Copy link
Member

@mkow mkow left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 1 of 1 files at r2, all commit messages.
Reviewable status: all files reviewed, 8 unresolved discussions, not enough approvals from maintainers (2 more required), not enough approvals from different teams (2 more required, approved so far: ), "fixup! " found in commit messages' one-liners (waiting on @kailun-qin)


a discussion (no related file):
@adarshan-intel: Please read our contributing guidelines and style guide, you're making a lot of small mistakes in your PRs.


-- commits line 3 at r2:

Suggestion:

pytorch: Set weights_only=False when loading models from a trusted source

-- commits line 5 at r2:
please wrap


-- commits line 5 at r2:
What's "direct unpickling" exactly? For me it sounds like the direct one would be with weights_only=False, not True.

Code quote:

direct unpickling

pytorch/pytorchexample.py line 7 at r2 (raw file):

import torch

# Load the model from a file (ensure this file is from a trusted source)

Suggestion:

# Load the model from a file. This file needs to be obtained from a trusted source, because it can contain code, not only data.

pytorch/pytorchexample.py line 9 at r2 (raw file):

# Load the model from a file (ensure this file is from a trusted source)
# Note: For PyTorch 2.6 and later, direct unpickling with weights_only=True (which was made the default)
# can fail if the model file contains more than just the weights and includes classes or functions that are not allowlisted.

please wrap

@adarshan-intel
Copy link
Contributor Author

-- commits line 5 at r2:

Previously, mkow (Michał Kowalczyk) wrote…

What's "direct unpickling" exactly? For me it sounds like the direct one would be with weights_only=False, not True.

Direct unpickling loads a serialized object (e.g., a PyTorch model) directly from a file using Python's pickle module, which can execute arbitrary code and pose a security risk if the file is untrusted.

In PyTorch, weights_only=True loads only the model's weights, avoiding potentially harmful code. However, if the model file contains more than just weights (e.g., custom classes or functions), weights_only=True might fail. Using weights_only=False loads the entire model, including any custom components, but is riskier if the source is untrusted.

Copy link
Contributor Author

@adarshan-intel adarshan-intel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewable status: all files reviewed, 8 unresolved discussions, not enough approvals from maintainers (2 more required), not enough approvals from different teams (2 more required, approved so far: ), "fixup! " found in commit messages' one-liners (waiting on @kailun-qin and @mkow)


-- commits line 5 at r2:

Previously, mkow (Michał Kowalczyk) wrote…

please wrap

Done.


pytorch/pytorchexample.py line 9 at r2 (raw file):

Previously, mkow (Michał Kowalczyk) wrote…

please wrap

Done.


-- commits line 3 at r2:
Done.


pytorch/pytorchexample.py line 7 at r2 (raw file):

import torch

# Load the model from a file (ensure this file is from a trusted source)

Done.

Copy link
Member

@mkow mkow left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 1 of 1 files at r3, all commit messages.
Reviewable status: all files reviewed, 5 unresolved discussions, not enough approvals from maintainers (2 more required), not enough approvals from different teams (2 more required, approved so far: ) (waiting on @adarshan-intel and @kailun-qin)


-- commits line 5 at r2:

Previously, adarshan-intel (Adarsh Anand) wrote…

Direct unpickling loads a serialized object (e.g., a PyTorch model) directly from a file using Python's pickle module, which can execute arbitrary code and pose a security risk if the file is untrusted.

In PyTorch, weights_only=True loads only the model's weights, avoiding potentially harmful code. However, if the model file contains more than just weights (e.g., custom classes or functions), weights_only=True might fail. Using weights_only=False loads the entire model, including any custom components, but is riskier if the source is untrusted.

Yeah, I know how unpickling works and the security risks. I was asking about that "direct unpickling" term you used, what does it mean exactly? What you described above is just how normal unpickling works, but saying "direct unpickling" implicates that there's also another way of unpickling (indirect?) - which I assumed is the one with weights_only=True, which doesn't run the Python unpickler but only parses weights from the file.
And here's the confusion - you say "direct unpickling with weights_only=True", which sounds self-contradictory.


pytorch/pytorchexample.py line 7 at r3 (raw file):

import torch

# Load the model from a file. This file needs to be obtained from a trusted 

trailing space

Copy link
Contributor Author

@adarshan-intel adarshan-intel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewable status: 0 of 1 files reviewed, 5 unresolved discussions, not enough approvals from maintainers (2 more required), not enough approvals from different teams (2 more required, approved so far: ), "fixup! " found in commit messages' one-liners (waiting on @kailun-qin and @mkow)


-- commits line 5 at r2:

Previously, mkow (Michał Kowalczyk) wrote…

Yeah, I know how unpickling works and the security risks. I was asking about that "direct unpickling" term you used, what does it mean exactly? What you described above is just how normal unpickling works, but saying "direct unpickling" implicates that there's also another way of unpickling (indirect?) - which I assumed is the one with weights_only=True, which doesn't run the Python unpickler but only parses weights from the file.
And here's the confusion - you say "direct unpickling with weights_only=True", which sounds self-contradictory.

So, when I mentioned “direct unpickling with weights_only=True,” it was indeed a bit misleading. We can go with normal unpickling, I don't think so there are any terms like direct or indirect pickling. Yeah and making from weights_only=True to weights_only=False does not mean going direct to indirect pickling.


pytorch/pytorchexample.py line 7 at r3 (raw file):

Previously, mkow (Michał Kowalczyk) wrote…

trailing space

Done.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

PyTorch 2.6.0 Release: Model Loading Fails with Unpickling Error
3 participants