Unveiling the Hidden Perils of the TorchScript Engine in PyTorch

Unveiling the Hidden Perils of the TorchScript Engine in PyTorch

Speakers: Jian Zhou and Lishu Song, Security Engineers, Alibaba Cloud

Conference: Black Hat USA 2025 — August 6-7, 2025, Mandalay Bay, Las Vegas

YouTube: https://www.youtube.com/watch?v=rZ0Jnu5gtqo

Reading Time: ~8 minutes

Type: Briefing

TL;DR

PyTorch's weights_only=True parameter — the standard fix for pickle-based RCE in ML model loading — does not actually prevent code execution when loading TorchScript (.pt) files, because the TorchScript engine processes a separate code path with over 2,000 built-in operators. Alibaba Cloud researchers discovered two operators (save and from_file) that enable arbitrary file write and read, confirmed RCE against VLLM and Hugging Face Transformers, and reported a heap overflow in the same code path. Both ecosystems have now patched, but hard-coded PyTorch version pins in downstream projects left users exposed long after fixes were available.

Introduction

PyTorch is the world's most widely used deep learning framework — the foundation of research pipelines, inference servers, and commercial AI products across every major cloud provider. Its .pt model file format is the default packaging mechanism for sharing and deploying models on platforms such as Hugging Face, making trust in model loading security critical to the entire AI supply chain.

For years, the known risk was pickle: PyTorch's default torch.load() function uses pickle deserialization, which can execute arbitrary code embedded in a malicious model file. The community-accepted fix was to set weights_only=True, which switches to a restricted unpickler with whitelist and blacklist controls on permitted operations. Alibaba Cloud engineers Jian Zhou and Lishu Song went one level deeper and found that weights_only=True is not the safe harbor everyone thought it was.

The Problem with Pickle and Why weights_only Seemed to Solve It

▶ Watch: Pickle RCE demonstration (02:00)

The security risk from pickle is well-established and even documented in Python's own official documentation. When torch.load() is called without weights_only=True, any model file can execute arbitrary Python via pickle's __reduce__ magic method. A two-line payload calling os.system('whoami') embedded in a .pt file demonstrates reliable RCE.

In 2021, a GitHub issue called for PyTorch to address the problem. By 2022, PyTorch released weights_only=True as the answer. The restricted unpickler it activates (WeightsOnlyUnpickler) adds whitelist and blacklist checks on the two most dangerous pickle opcodes: GLOBAL (which loads a Python class) and REDUCE (which calls it). The community treated this as a solved problem — and the fix became the standard patch for downstream CVEs. When VLLM was found using torch.load without weights_only=True, the patch was simply adding the parameter. When Hugging Face Transformers had the same issue, same fix.

Discovering the TorchScript Code Path

▶ Watch: Finding torch._jit_internal.load (08:00)

Zhou and Song conducted a comprehensive audit of the full weights_only=True loading workflow rather than just the pickle layer. During step-through debugging, they observed a call to torch._jit_internal.load — a function they did not initially recognize.

This is the TorchScript engine. TorchScript is PyTorch's mechanism for exporting Python-defined models to environments without a Python interpreter, such as C++ deployments and mobile applications. A .pt file saved via torch.jit.save() is a zip archive containing:

• data.pkl — serialized object structure (the model's R-values in pickle format)
• constants.pkl — tensor constants in pickle format
• code/ directory — Python source code regenerated from the compiled IR (Intermediate Representation) graph

When weights_only=True is set but the file being loaded is a TorchScript model, PyTorch invokes the TorchScript deserializer (deserialize → reader_archive), which processes the data.pkl and constants.pkl through pickle — but also reads and compiles the source code from the code/ directory. Crucially, this code is not subject to the WeightsOnlyUnpickler restrictions.

2,000 Operators and Two Dangerous Ones

▶ Watch: Operator table discovery (20:01)

TorchScript code executes via an IR graph where each node contains a gen, output, and opcode. One opcode, OP, performs a lookup in an operator table that maps strings to function pointers. Zhou patched PyTorch to enumerate all registered operators and found over 2,000 entries.

A code audit of those operators revealed two of particular interest:

• save (exposed via torch.save in TorchScript context): Saves a specified value as a model to a specified file path — enabling arbitrary file write.
• from_file (exposed via torch.from_file): Reads a file path and converts its contents to a tensor, from which the raw bytes can be recovered — enabling arbitrary file read.

The save operator's write capability is constrained by its permission model (files are written with 644 permissions), which complicates overwriting certain system files. The cron daemon, for example, requires 600 permissions — a direct overwrite attempt fails due to the permission mismatch. However, the researchers found working exploitation paths.

RCE Confirmed Against VLLM and Hugging Face Transformers

▶ Watch: VLLM exploitation demo (28:39)

To reach these operators from within TorchScript, the researchers needed to understand how they are registered and exposed. The get_builtin_table function registers object member functions to the builtin operator table with a name prefix transformation. During TorchScript compilation, to_sugar_val checks whether a Python object exists in that table and, if so, emits a builtin function into the IR graph. The save operator is accessible via torch (the torch Python object), making torch.save() and torch.from_file() directly callable in crafted TorchScript source code.

VLLM exploitation: VLLM v0.7.3 patched its original torch.load-without-weights_only CVE by adding weights_only=True. The patch appeared effective but was not. The loaded model's items() function is called after loading; by crafting a TorchScript model that exports a method named items decorated with @torch.jit.export, the researchers triggered the malicious TorchScript code path through the normal VLLM inference workflow. Additionally, VLLM hard-coded a specific PyTorch version in its requirements_cpu.txt and requirements_gpu.txt, meaning users could not update PyTorch even after PyTorch released a patch.

Hugging Face Transformers exploitation: Transformers v4.51.3, which defaults to weights_only=True, loads .bin format model files by calling torch._jit_internal.load. By crafting a malicious pytorch_model.bin in a local repository and loading it via Transformers' standard model-loading API, the researchers triggered RCE. The equivalent spoofed method name in this path is keys — a standard Python dict method. The patch in both cases was simply requiring PyTorch ≥ 2.6.

Heap Overflow in the Same Code Path

▶ Watch: Memory corruption disclosure (26:39)

Beyond the logic vulnerability, the researchers also identified memory corruption bugs in the TorchScript loading path. Loading a crafted model file triggered a heap overflow that leaked both heap addresses and library base addresses — the prerequisites for bypassing ASLR in a memory corruption exploit chain. Due to time constraints, full exploitation details were not presented, but the researchers indicated plans to share them in future work.

Defense Recommendations

PyTorch's fix prevents TorchScript from being invoked when weights_only=True is set, effectively closing the disclosed path. Recommendations for the broader ecosystem:

• Update PyTorch to ≥ 2.6 immediately. Do not accept hard-coded version pins in downstream libraries as an excuse not to update.
• Prefer safetensors format for model storage and distribution. The safetensors format was designed specifically to avoid pickle and code execution during model loading.
• Scan model repositories for malicious content before loading. Hugging Face has implemented server-side malware scanning, but local or private repositories require independent validation.
• Never load models from untrusted sources, regardless of the weights_only parameter. Load models in isolated environments (containers, VMs) when the source cannot be verified.
• Audit downstream libraries that call torch.load or torch._jit_internal.load and verify they are using the patched PyTorch version.

Notable Quotes

"The safe harbor we once thought was actually hostile waters." — Jian Zhou ▶ 28:39

"Can we blindly trust the security of weights_only? The answer is no." — Jian Zhou ▶ 04:00

"The entire AI ecosystem is just like a house, and PyTorch is the fundamental component. If there is an issue with PyTorch, the security of the entire structure is compromised." — Jian Zhou ▶ 26:39

Key Takeaways

• PyTorch's weights_only=True does not prevent RCE when loading TorchScript (.pt) files, because the TorchScript engine has its own code execution path with over 2,000 built-in operators.
• Two operators — torch.save (arbitrary file write) and torch.from_file (arbitrary file read) — can be called from crafted TorchScript source code embedded in a model file.
• VLLM and Hugging Face Transformers were both vulnerable even after applying their "patched" weights_only=True mitigations; hard-coded PyTorch version pins in VLLM's requirements files further delayed user protection.
• A heap overflow in the TorchScript loading path additionally leaks memory addresses, enabling potential memory corruption exploit chains.
• The fix is PyTorch ≥ 2.6, which blocks TorchScript from running under weights_only=True; the recommended long-term solution is migrating to the safetensors model format.

Slides for this talk were not listed as available in the session bundle.