AI
The False Promise of weights_only
PyTorch used pickle to load models by default. Pickle executes arbitrary code: an attacker crafts a model that calls os.system through the REDUCE opcode. In 2021, someone flagged this on GitHub. In 2022, PyTorch shipped weights_only=True, a restricted pickle that whitelists load_global and load_reduce. Projects like vLLM got CVEs for loading without the parameter, then patched by adding it. The PyTorch docs called weights_only=True safe. The real question is whether that safety claim holds up under scrutiny.
TorchScript Internals: Compilation, Serialization, and Execution
TorchScript compiles Python model code to an IR graph, then serializes it as a .pt file (a ZIP archive). Inside: data.pkl holds the model object, constants.pkl holds tensor constants, and code/ holds source strings. Loading calls torch.jit.load, which unpickles both .pkl files and parses the source into IR. Execution iterates the IR graph, dispatching each opcode through a switch table. One opcode, OP, routes calls through an operator table with over 2,000 registered functions, none of them gated by weights_only.
Arbitrary File Read and Write via the Operator Table
Two operators in that table stood out: save and from_file. The save operator writes a model to any path the TorchScript specifies. The from_file operator reads any file and returns its contents as a tensor. Both are callable from a scripted model via torch.save or torch.from_file. Writing to ~/.ssh/authorized_keys failed: torch.save creates files with 644 permissions, but SSH requires 600. A different target worked. The TorchScript arbitrary file write PoC (25:58) shows the result:
“You can see we get a CE.” (Ji’an Zhou)
Downstream Impact on vLLM and Hugging Face Transformers
vLLM had already received a CVE for loading without weights_only=True and patched it by adding the parameter. But vLLM’s requirements.txt hardcoded a specific PyTorch version for both CPU and GPU installs, so users couldn’t update to a fixed PyTorch even after one shipped. Hugging Face Transformers set weights_only=True by default and still fell to the same torch.jit.load path. The fix for both required PyTorch >= 2.6, the first version where torch.jit.load is blocked when weights_only is enabled.
Notable Quotes
can we blendly trust the security of v Ji’an Zhou · ▶ 4:50
You can see we get a CE. Ji’an Zhou · ▶ 26:39
because PyTorch document just claim it safe Ji’an Zhou · ▶ 33:26
Key Takeaways
- weights_only=True blocks pickle deserialization but leaves torch.jit.load fully exposed.
- TorchScript’s operator table includes save and from_file, giving arbitrary file read/write from a .pt model.
- vLLM and Hugging Face Transformers both remained vulnerable after patching the original pickle CVEs.