Unveiling the Hidden Perils of the TorchScript Engine in PyTorch

Overview: Security researchers Guorui (Gero) and Liz Ross from Alibaba Cloud analyze the fundamental security limitations of PyTorch's weights_only=True loading parameter, revealing that upstream protections against Python pickle code execution did not prevent remote code execution. While widely adopted by the AI community and downstream frameworks to neutralize malicious __reduce__ pickle payloads, the authors discovered that torch.load() continued to pass TorchScript model archives to the torch.jit._load execution engine. By exploiting vulnerable built-in TorchScript operators and method-export behaviors, the researchers achieved arbitrary file read/write, memory corruption, and remote code execution against major LLM serving frameworks including vLLM and Hugging Face Transformers.

Key Topics: - Limitations of restricted unpickling and PyTorch's weights_only=True parameter - Architecture of the TorchScript JIT compiler, Intermediate Representation (IR) graphs, and operator dispatch - Exploitation of built-in operators (torch.save, torch.from_file) and deserialization heap overflows - Downstream weaponization against vLLM and Hugging Face Transformers via method spoofing (items() and keys())

Technical Details: To mitigate classic Python pickle deserialization attacks (where malicious models exploit GLOBAL and REDUCE opcodes to execute os.system), PyTorch introduced weights_only=True in torch.load(). This mechanism employs a restricted unpickler enforcing an opcode and class allowlist. However, auditing the complete torch.load() workflow revealed an overlooked execution branch: when handling serialized TorchScript archives (ZIP containers containing data.pkl, constants.pkl, and IR code directories), torch.load forwarded the payload directly to torch.jit._load.load(). TorchScript compiles Python code into intermediate representation (IR) graphs composed of nodes and opcodes executed via an internal switch table that interfaces with over 2,000 registered C++ operators managed by RegisterOperators.

Auditing these internal operators uncovered critical built-in primitives that could be called directly from TorchScript IR without security boundaries: torch.save provided arbitrary file-writing capabilities, while torch.from_file enabled arbitrary file-reading by mapping file contents directly into tensors. The researchers also identified heap overflow vulnerabilities within the TorchScript deserialization parser, allowing attackers to corrupt memory and leak heap and library base addresses. To weaponize this vector against real-world applications, the researchers examined downstream frameworks that integrated weights_only=True. In vLLM (v0.7.3), loaded models are inspected by invoking model.items() rather than forward(). Because TorchScript applies a lazy export strategy for functions unless explicitly decorated with @torch.jit.export, the authors decorated a spoofed items() method to ensure it was exported into the IR graph, successfully triggering the malicious operator and achieving RCE during model loading. Similarly, in Hugging Face Transformers (v4.51.3), model ingestion routines invoked model.keys(), which the researchers exploited by exporting a spoofed keys() function.

Key Takeaways: Relying on high-level parameter flags like weights_only=True can create a false sense of security if underlying framework parsers dynamically route untrusted data into complex compilation or JIT execution engines. Following responsible disclosure, PyTorch patched the vulnerability in version 2.6 by strictly prohibiting the instantiation of the TorchScript engine when weights_only=True is asserted. Downstream AI serving frameworks must avoid hardcoding and pinning vulnerable PyTorch dependencies, and organizations should transition AI pipelines away from executable formats like .pt and .bin toward static, safe serialization formats such as Safetensors.

(AI summary — Gemini Flash)