`_set_obj_state` applies unrestricted setattr from checkpoint state under weights_only=True
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
## Summary
When loading a checkpoint with `torch.load(..., weights_only=True)`, `_set_obj_state()` in `torch/_utils.py` calls `setattr(obj, k, v)` using keys taken straight from the pickle stream, with no filtering. A crafted state dict can set dunder attributes like `__class__` on a loaded object, which silently changes its type (for example `nn.Parameter` into a plain `Tensor`) with no error or warning.
To be clear, this is not a security vulnerability. `weights_only=True` guarantees safe deserialization (no RCE, no out-of-bounds writes), and the values here still go through the restricted unpickler. This is a hardening request. A normal `torch.save` never emits these keys, so rejecting them is cheap and closes off a class of surprising, hard-to-debug behavior from malformed or hand-crafted checkpoints.
## Location
- File: `torch/_utils.py`, function `_set_obj_state()`
- Callers: `_rebuild_parameter_with_state` and `torch._tensor._rebuild_from_type_v2`, both on the restricted unpickler's allowlist
```python
def _set_obj_state(obj, state):
...
for k, v in dict_state.items():
setattr(obj, k, v) # k comes from the pickle stream, unfiltered
```
## What happens
The restricted unpickler validates a lot of operations. `GLOBAL` is checked against an allowlist, `REDUCE` only permits allowed functions, and `SETITEM`/`SETITEMS` is restricted to `dict`, `OrderedDict`, and `Counter`. What it doesn't restrict is which attribute names can be set on objects returned by allowed functions. `_set_obj_state` is the one path where a key from the stream reaches `setattr` directly.
Since `Parameter` inherits directly from `Tensor` (same C-level layout), CPython allows a `__class__` reassignment, so a state dict of `{"__class__": torch.Tensor}` downgrades a `Parameter` to a plain `Tensor`. `isinstance(x, nn.Parameter)` then returns `False`, which quietly breaks optimizers, gradient tracking, and model-surgery tools downstream, with no signal at load time.
The same mechanism also lets you inject arbitrary attributes (for example `{"_admin": True}`) into `__dict__`, or set `__reduce__` or `__setattr__` to non-callable values, which then raises `TypeError` later on re-serialization or attribute assignment.
## Reproduction
Reproduced on PyTorch 2.13.0+cu130, Python 3.13, Colab (Tesla T4). Self-contained, no downloads needed.
```python
import struct, tempfile, zipfile, os
import torch
PROTO=b'\x80'; STOP=b'.'; GLOBAL=b'c'; REDUCE=b'R'
MARK=b'('; TUPLE=b't'; EMPTY_DICT=b'}'; EMPTY_TUPLE=b')'
SETITEM=b's'; SETITEMS=b'u'; BINPUT=b'q'
BINUNICODE=b'X'; NEWFALSE=b'\x89'; NEWTRUE=b'\x88'
BININT1=b'K'; BINPERSID=b'Q'
def ustr(s):
b = s.encode('utf-8')
return BINUNICODE + struct.pack('
Contributor guide
Assessment
This issue has not been assessed yet.