Unveiling the Perils of the TorchScript Engine in PyTorch

Unveiling the Perils of the TorchScript Engine in PyTorch

Speakers: Ji'an Zhou, Lishuo Song

Conference: DEF CON 33

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

Slides: https://media.defcon.org/DEF%20CON%2033/DEF%20CON%2033%20presentations/Ji%27an%20Zhou%20%26%20Lishuo%20Song%20-%20Safe%20Harbor%20or%20Hostile%20Waters%20Unveiling%20the%20Hidden%20Perils%20of%20the%20TorchScript%20Engine%20in%20PyTorch.pdf

Overview

This DEF CON 33 talk—titled in full "Safe Harbor or Hostile Waters: Unveiling the Hidden Perils of the TorchScript Engine in PyTorch"—presents a systematic security analysis of PyTorch's TorchScript engine, a component designed to compile Python-based machine learning models into a serializable, deployable representation. Security engineers Ji'an Zhou and Lishuo Song from Alibaba Cloud discovered that TorchScript, which is widely used to distribute pre-trained models, contains serious vulnerabilities that allow arbitrary code execution when a malicious model file is loaded—even in environments where the use of Python's unsafe pickle serialization format has been explicitly abandoned.

The talk is significant because it attacks the "safe harbor" narrative around TorchScript: when PyTorch deprecated the legacy torch.load pickle-based loading in favor of safer alternatives, TorchScript was often positioned as a secure serialization path. The researchers demonstrate this assumption is wrong.

Background

▶ Watch: Background: PyTorch TorchScript as a model serialization and execution engine (2:14)

PyTorch is the dominant machine learning framework in research and production. According to benchmarks cited in the talk (referencing a "Papers with Code" survey), PyTorch surpassed TensorFlow as the most widely used ML framework for deep learning research and has substantial production deployment.

The serialization problem. ML models need to be saved and loaded—for deployment, transfer learning, and sharing via model repositories like Hugging Face. PyTorch's original serialization mechanism used Python's pickle module, which is notoriously unsafe: loading a pickled file from an untrusted source allows arbitrary code execution at deserialization time. This was widely publicized and led to many supply chain attacks through model repositories.

TorchScript as the "safe" alternative. TorchScript is PyTorch's compiled representation format. A Python model is converted to TorchScript using torch.jit.script() or torch.jit.trace(), producing a .pt file. Unlike pickle-based .pt files, TorchScript files use a different internal format—they serialize a computation graph rather than Python objects. The prevailing assumption was that this made TorchScript loading safe for untrusted model files.

The research question: Zhou and Song set out to determine whether this assumption holds—whether the TorchScript engine, specifically the parser and interpreter that execute TorchScript code at load time, is free from exploitable vulnerabilities.

Key Findings

▶ Watch: Research methodology: breaking out of conventional vulnerability search patterns (11:20)

TorchScript is not a safe serialization format for untrusted models. The researchers discovered multiple classes of vulnerabilities in the TorchScript engine:

1. Arbitrary file read vulnerabilities. The TorchScript VM includes built-in operations that can access the filesystem. By crafting a malicious TorchScript model, an attacker can trigger file reads from arbitrary paths on the loading system, enabling credential theft and data exfiltration.

1. Memory corruption bugs in the TorchScript parser. The TorchScript parser, written in C++, contains type confusion and use-after-free vulnerabilities that can be triggered by specially crafted TorchScript source representations embedded in a model file. These bugs are reachable through the normal model loading code path.

1. Server-Side Request Forgery (SSRF) via TorchScript operations. TorchScript includes tensor operations that can invoke network operations. A malicious model can trigger outbound network requests from the loading server, enabling SSRF attacks against internal cloud metadata services (such as AWS IMDSv1).

1. Sandbox escape in model serving platforms. Model serving platforms (such as TorchServe) that accept user-provided models and load them in a sandboxed environment are vulnerable to escaping the sandbox through these TorchScript bugs.

Impact. The combination of these vulnerabilities means that any organization that accepts model uploads, serves models from a public repository, or loads TorchScript models from untrusted sources is vulnerable to compromise. This directly threatens AI/ML pipelines, model repositories, and MLOps platforms.

CVE assignments. The researchers reported their findings to the PyTorch security team and multiple CVEs were issued for the discovered vulnerabilities prior to the talk.

Technical Deep Dive

▶ Watch: TorchScript VM internals: JIT compiler and serialization attack surface (15:45)

TorchScript Architecture

TorchScript has two main components:

1. The frontend (script/trace): Converts Python model code to an intermediate representation (IR)—a computation graph with typed nodes. The torch.jit.script path performs static analysis of Python code; torch.jit.trace records operations from a forward pass.

1. The TorchScript VM: A C++-based virtual machine that executes the IR. The VM includes a bytecode interpreter, a type system, and built-in operator implementations (arithmetic, tensor operations, I/O, etc.).

When a .pt file containing a TorchScript model is loaded with torch.jit.load(), the VM deserializes and compiles the embedded IR, then instantiates the model. The critical security observation is that the IR execution happens automatically at load time—there is no "safe inspect first" option.

Vulnerability Classes

Type Confusion in the IR Parser

The TorchScript IR uses a statically typed node graph. However, certain optimizations in the C++ parser perform type inference and can be confused by crafted IR that presents a type mismatch between the node declaration and its usage. This can lead to memory corruption—specifically, treating a pointer to one object type as a pointer to another (type confusion), which can be leveraged for read/write primitives.

Use-After-Free in Optimization Passes

The TorchScript compiler applies several optimization passes (constant folding, dead code elimination, etc.) over the IR graph. These passes manipulate the graph's node list. A crafted IR graph can trigger a node deletion followed by a use of the same node pointer in the same pass, resulting in a use-after-free condition. In a C++ context, this is a classic memory safety bug that can enable arbitrary code execution.

Built-in Filesystem Operations

TorchScript exposes Python-like built-in types, including string manipulation and file I/O operations inherited from the Python embedding. The researchers identified specific TorchScript built-ins that invoke filesystem syscalls without sanitization. An attacker can write a model that, when loaded, reads /etc/passwd, cloud metadata from 169.254.169.254, or private key files from well-known paths.

SSRF via HTTP Builtins

Certain TorchScript operations (related to custom operator loading and remote resource fetching in model components) can trigger outbound HTTP requests. This enables SSRF, particularly dangerous in cloud environments where instance metadata is accessible at a fixed IP.

Attack Scenario: Model Repository Poisoning

The practical attack scenario the researchers present is a supply chain attack:

1. Attacker uploads a malicious TorchScript model to Hugging Face, PyPI, or an internal model registry.
2. A data scientist or MLOps pipeline downloads and loads the model using torch.jit.load().
3. The model executes the malicious IR at load time, exfiltrating credentials or executing a reverse shell payload.
4. The attack is silent—the model may still function correctly for its stated purpose, making detection difficult.

Demo / Proof of Concept

▶ Watch: Vulnerability class 1: arbitrary file read through TorchScript builtins (24:44)

The talk was presented via pre-recorded video (the speakers were unable to attend in person due to visa issues), but it included:

• PoC model files demonstrating arbitrary file read, showing /etc/passwd contents extracted from the loading server.
• SSRF demonstration showing a crafted model triggering an HTTP request to AWS IMDSv1 (169.254.169.254/latest/meta-data/) and returning the instance IAM role credentials.
• Memory corruption PoC demonstrating crash conditions triggered by the type confusion bug, with analysis of the resulting crash state.

Full proof-of-concept code was made available to the PyTorch security team prior to disclosure.

Defensive Implications

▶ Watch: CVE selection: choosing most impactful bugs for responsible disclosure (33:50)

For ML engineers and data scientists:

• Do not load TorchScript model files from untrusted sources. The safety of TorchScript relative to pickle has been overstated.
• Prefer hash-verified model loading: always verify the SHA256 hash of a model file against a trusted manifest before loading it.
• If loading models from external repositories, use model cards and provenance attestation (e.g., Hugging Face's model signing initiative) to verify the model's origin.

For MLOps and model serving platforms:

• Sandbox model loading in a separate process with minimal filesystem and network access (seccomp-BPF, namespaces, or a dedicated VM).
• Treat user-submitted models as untrusted code, not as data files.
• Implement network egress controls that prevent outbound connections from the model loading environment.
• Audit your serving infrastructure for exposure to cloud metadata endpoints.

For PyTorch maintainers and the broader framework ecosystem:

• The TorchScript VM should be treated as an interpreter of untrusted code, with appropriate sandboxing.
• A "safe mode" for model loading that disables filesystem and network operations would substantially reduce the attack surface.
• Expanding fuzzing coverage of the TorchScript parser and VM (analogous to the work presented in the dc33-087 Orion talk) would proactively find additional memory safety bugs.

Key Takeaways

1. TorchScript is not a safe alternative to pickle for loading untrusted models—it has its own code execution surface through the TorchScript VM.
2. Multiple vulnerability classes (type confusion, use-after-free, filesystem access, SSRF) exist in the TorchScript engine and are reachable through normal model loading.
3. ML supply chain attacks via malicious model files are a realistic, high-impact threat to organizations using PyTorch at scale.
4. The assumption that "compiled model formats are safe to load from untrusted sources" is incorrect for TorchScript and likely for other compiled ML representations as well.
5. Defenders should apply the same scrutiny to model files that they apply to code artifacts—hash verification, provenance tracking, and sandboxed execution.

About the Speaker(s)

▶ Watch: Demo: renaming model function to trigger RCE on model load (38:05)

Ji'an Zhou and Lishuo Song are security engineers at Alibaba Cloud. Their research focuses on the intersection of machine learning frameworks and systems security. The work on TorchScript was conducted as part of Alibaba Cloud's security research program, which investigates security vulnerabilities in widely adopted open-source infrastructure. The speakers were unable to attend DEF CON 33 in person due to visa application issues and presented via pre-recorded video, with the in-person audience able to submit questions for follow-up.