ghstack-source-id:66a70e7Pull Requestresolved:#129509

ghstack-source-id:7d6ee40Pull Requestresolved:#129509

ghstack-source-id:7d6ee40Pull Requestresolved:pytorch#129509

* Fix allowlisting of builtins for weights_only unpickler (#129244)Since we use [`DEFAULT_PROTOCOL=2`](https://github.com/pytorch/pytorch/blob/main/torch/serialization.py#L62), some functions/classes that were renamed from python 2-->3 will be pickled with their python2 name. This PR ensures that when a mod `GLOBAL <python2_mod>.<python2_name> ` is encountered, [following the strategy used by pickle](https://github.com/python/cpython/blob/main/Lib/pickle.py#L1590C13-L1593C63) it is properly mapped to `<python3_mod>.<python3_name>`.This fix ensures that `add_safe_globals` works properly for such functions/classes (i.e. users will allowlist the python3 func and the weights_only unpickler will do the appropriate translation when checking whether a class was allowlisted).An example is as follows:`__builtin__` was named to `builtins`, see the [release notes for Python 3.0](https://docs.python.org/3/whatsnew/3.0.html)> Renamed module `__builtin__` to [`builtins`](https://docs.python.org/3/library/builtins.html#module-builtins) (removing the underscores, adding an ‘s’). The __builtins__ variable found in most global namespaces is unchanged. To modify a builtin, you should use [builtins](https://docs.python.org/3/library/builtins.html#module-builtins), not `__builtins__`!However, since we use [`DEFAULT_PROTOCOL=2`](https://github.com/pytorch/pytorch/blob/main/torch/serialization.py#L62), builtins will be pickled with their module string as `__builtin__`.```python>>> import pickle>>> import pickletools>>> print.__module__'builtins'>>> with open('print.pkl', 'wb') as f:>>>      pickle.dump(print, f, protocol=2) # 2 because this is the default protocol used by pytorch>>> with open('print.pkl', 'rb') as f:>>>     pickletools.dis(f)0: \x80 PROTO      22: c    GLOBAL     '__builtin__ print' # pickle saves the module string as __builtin__ !!! :(21: q    BINPUT     023: .    STOP```Pull Requestresolved:#129244Approved by:https://github.com/albanD* Allow BUILD/NEWOBJ instruction for items added via torch.serialization.add_safe_globals (#129251)Previously, allowlisting functions/classes via `torch.serialization.add_safe_globals(obj)` for the `weights_only` Unpickler had the following effect:- For a [`GLOBAL`](https://github.com/python/cpython/blob/3.12/Lib/pickletools.py#L1926-L1939) instruction, `GLOBAL obj.__module__ obj.__name__` would be allowed and translated back to obj to be pushed back to the stack.- For a [`REDUCE`](https://github.com/python/cpython/blob/3.12/Lib/pickletools.py#L1926-L1982) instruction where we expect the stack to contain `func` and `args`, `func` is allowed if it was added via `add_safe_globals`However, it did not have an effect on `BUILD` and `NEWOBJ` instructionsSome classes may be rebuilt via [`NEWOBJ`](https://github.com/python/cpython/blob/3.12/Lib/pickletools.py#L2091-L2104) instruction, which indicates that their constructor should be used to rebuild the class.Further, a [`BUILD`](https://github.com/python/cpython/blob/3.12/Lib/pickletools.py#L1984-L2007) instruction might be used if an object's `__reduce__`/`__reduce_ex__` returns a non-None value for `state`. Which indicates a `__setstate__` or `__dict__.update`.**This PR makes sure that adding objects to the allowlist will also allow `NEWOBJ` and `BUILD` instructions for them.**In particular, the update for `NEWOBJ` should unblock allowlisting of [`ScaledMMConfig`](https://github.com/pytorch-labs/float8_experimental/blob/d4ade877dff327ea7f51e91f7cc218ae956e8cfd/float8_experimental/float8_tensor.py#L26-L30) in float8_experimental@drisspgPull Requestresolved:#129251Approved by:https://github.com/albanDghstack dependencies:#129244* Remove dependency on private _compat_pickle in CPythonghstack-source-id:7d6ee40Pull Requestresolved:#129509

As@vmoens pointed out, the current error message does not make the "either/or" between setting `weights_only=False` and using `add_safe_globals` clear enough, and should print the code for the user to call `add_safe_globals`New formatting looks like suchIn the case that `add_safe_globals` can be used```python>>> import torch>>> from torch.testing._internal.two_tensor import TwoTensor>>> torch.save(TwoTensor(torch.randn(2), torch.randn(2)), "two_tensor.pt")>>> torch.load("two_tensor.pt", weights_only=True)Traceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1225, in load    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options        (1) 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.        (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.        WeightsUnpickler error: Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor was not an allowed global by default. Please use `torch.serialization.add_safe_globals([TwoTensor])` to allowlist this global if you trust this class/function.Check the documentation of torch.load to learn more about types accepted by default with weights_onlyhttps://pytorch.org/docs/stable/generated/torch.load.html.```For other issues (unsupported bytecode)```python>>> import torch>>> t = torch.randn(2, 3)>>> torch.save(t, "protocol_5.pt", pickle_protocol=5)>>> torch.load("protocol_5.pt", weights_only=True)/data/users/mg1998/pytorch/torch/_weights_only_unpickler.py:359: UserWarning: Detected pickle protocol 5 in the checkpoint, which was not the default pickle protocol used by `torch.load` (2). The weights_only Unpickler might not support all instructions implemented by this protocol, please file an issue for adding support if you encounter this.  warnings.warn(Traceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1225, in load    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None_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: Unsupported operand 149Check the documentation of torch.load to learn more about types accepted by default with weights_onlyhttps://pytorch.org/docs/stable/generated/torch.load.html.```Old formatting would have been like:```pythonTraceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1203, in load    raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None_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 get the file from a trusted source. Alternatively, to load with `weights_only` please check the recommended steps in the following error message. WeightsUnpickler error: Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor was not an allowed global by default. Please use `torch.serialization.add_safe_globals` to allowlist this global if you trust this class/function.```Pull Requestresolved:#129705Approved by:https://github.com/albanD,https://github.com/vmoensghstack dependencies:#129239,#129396,#129509

As@vmoens pointed out, the current error message does not make the "either/or" between setting `weights_only=False` and using `add_safe_globals` clear enough, and should print the code for the user to call `add_safe_globals`New formatting looks like suchIn the case that `add_safe_globals` can be used```python>>> import torch>>> from torch.testing._internal.two_tensor import TwoTensor>>> torch.save(TwoTensor(torch.randn(2), torch.randn(2)), "two_tensor.pt")>>> torch.load("two_tensor.pt", weights_only=True)Traceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1225, in load    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options        (1) 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.        (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.        WeightsUnpickler error: Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor was not an allowed global by default. Please use `torch.serialization.add_safe_globals([TwoTensor])` to allowlist this global if you trust this class/function.Check the documentation of torch.load to learn more about types accepted by default with weights_onlyhttps://pytorch.org/docs/stable/generated/torch.load.html.```For other issues (unsupported bytecode)```python>>> import torch>>> t = torch.randn(2, 3)>>> torch.save(t, "protocol_5.pt", pickle_protocol=5)>>> torch.load("protocol_5.pt", weights_only=True)/data/users/mg1998/pytorch/torch/_weights_only_unpickler.py:359: UserWarning: Detected pickle protocol 5 in the checkpoint, which was not the default pickle protocol used by `torch.load` (2). The weights_only Unpickler might not support all instructions implemented by this protocol, please file an issue for adding support if you encounter this.  warnings.warn(Traceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1225, in load    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None_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: Unsupported operand 149Check the documentation of torch.load to learn more about types accepted by default with weights_onlyhttps://pytorch.org/docs/stable/generated/torch.load.html.```Old formatting would have been like:```pythonTraceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1203, in load    raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None_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 get the file from a trusted source. Alternatively, to load with `weights_only` please check the recommended steps in the following error message. WeightsUnpickler error: Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor was not an allowed global by default. Please use `torch.serialization.add_safe_globals` to allowlist this global if you trust this class/function.```Pull Requestresolved:#129705Approved by:https://github.com/albanD,https://github.com/vmoensghstack dependencies:#129239,#129396,#129509(cherry picked from commit45f3e20)

* Improve error message for weights_only load (#129705)As@vmoens pointed out, the current error message does not make the "either/or" between setting `weights_only=False` and using `add_safe_globals` clear enough, and should print the code for the user to call `add_safe_globals`New formatting looks like suchIn the case that `add_safe_globals` can be used```python>>> import torch>>> from torch.testing._internal.two_tensor import TwoTensor>>> torch.save(TwoTensor(torch.randn(2), torch.randn(2)), "two_tensor.pt")>>> torch.load("two_tensor.pt", weights_only=True)Traceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1225, in load    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options        (1) 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.        (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.        WeightsUnpickler error: Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor was not an allowed global by default. Please use `torch.serialization.add_safe_globals([TwoTensor])` to allowlist this global if you trust this class/function.Check the documentation of torch.load to learn more about types accepted by default with weights_onlyhttps://pytorch.org/docs/stable/generated/torch.load.html.```For other issues (unsupported bytecode)```python>>> import torch>>> t = torch.randn(2, 3)>>> torch.save(t, "protocol_5.pt", pickle_protocol=5)>>> torch.load("protocol_5.pt", weights_only=True)/data/users/mg1998/pytorch/torch/_weights_only_unpickler.py:359: UserWarning: Detected pickle protocol 5 in the checkpoint, which was not the default pickle protocol used by `torch.load` (2). The weights_only Unpickler might not support all instructions implemented by this protocol, please file an issue for adding support if you encounter this.  warnings.warn(Traceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1225, in load    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None_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: Unsupported operand 149Check the documentation of torch.load to learn more about types accepted by default with weights_onlyhttps://pytorch.org/docs/stable/generated/torch.load.html.```Old formatting would have been like:```pythonTraceback (most recent call last):  File "<stdin>", line 1, in <module>  File "/data/users/mg1998/pytorch/torch/serialization.py", line 1203, in load    raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None_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 get the file from a trusted source. Alternatively, to load with `weights_only` please check the recommended steps in the following error message. WeightsUnpickler error: Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor was not an allowed global by default. Please use `torch.serialization.add_safe_globals` to allowlist this global if you trust this class/function.```Pull Requestresolved:#129705Approved by:https://github.com/albanD,https://github.com/vmoensghstack dependencies:#129239,#129396,#129509(cherry picked from commit45f3e20)* Fix pickle import when rebase onto release/2.4* Update torch/serialization.pyfix bad rebase again---------Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>