frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Clarify error messages for NEWOBJ and BUILD in weights_only unpickler (#134346)

Browse Source 
Clarify that `add_safe_globals` will allow types for these instructions

Some types do not appear as `GLOBAL` and are only caught in `BUILD`, example from hf slack is `numpy.dtypes.UInt32DType`

```python
import torch
import numpy as np
from tempfile import TemporaryDirectory
from pathlib import Path
from codecs import encode

torch.serialization.add_safe_globals([encode, np.dtype, np.core.multiarray._reconstruct, np.ndarray])

with TemporaryDirectory() as tempdir:
    p = Path(tempdir)
    r2 = np.random.get_state()
    torch.save(r2, p / "r2.pkl")
    torch.load(p / "r2.pkl", weights_only=True)
```

Yields (error comes from BUILD)
```
UnpicklingError: Weights only load failed. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
 Please file an issue with the following so that we can make `weights_only=True` compatible with your use case: WeightsUnpickler error: Can only build Tensor, parameter or OrderedDict objects, but got <class 'numpy.dtypes.UInt32DType'>
```

The reasoning is that `numpy.dtypes.UInt32DType` is constructed via `REDUCE` with `func =<class 'numpy.dtype'>` and `args= ('u4', False, True)`, clarify the error message that doing `add_safe_globals` on these will also allow them

After this PR error message becomes

```
_pickle.UnpicklingError: Weights only load failed. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
Please file an issue with the following so that we can make `weights_only=True` compatible with your use case: WeightsUnpickler error: Can only build Tensor, Parameter, OrderedDict or types allowlisted via `add_safe_globals`, but got <class 'numpy.dtypes.UInt32DType'>
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/134346
Approved by: https://github.com/albanD
This commit is contained in:
Mikayla Gawarecki
2024-08-26 14:28:18 -07:00
committed by PyTorch MergeBot
parent 2ac710e667
commit 2033934ff8
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/_weights_only_unpickler.py
View File

@ -261,7 +261,8 @@ class Unpickler:
self.append(cls.__new__(cls, *args)) self.append(cls.__new__(cls, *args))
else: else:
raise UnpicklingError( raise UnpicklingError(
f"Trying to instantiate unsupported class {cls}" "Can only create new object for nn.Parameter or classes allowlisted "
f"via `add_safe_globals` but got {cls}"
) )
elif key[0] == REDUCE[0]: elif key[0] == REDUCE[0]:
args = self.stack.pop() args = self.stack.pop()
@ -291,7 +292,8 @@ class Unpickler:
inst.__dict__.update(state) inst.__dict__.update(state)
else: else:
raise UnpicklingError( raise UnpicklingError(
f"Can only build Tensor, parameter or OrderedDict objects, but got {type(inst)}" "Can only build Tensor, Parameter, OrderedDict or types allowlisted "
f"via `add_safe_globals`, but got {type(inst)}"
) )
# Stack manipulation # Stack manipulation
elif key[0] == APPEND[0]: elif key[0] == APPEND[0]: