Skip to content

Unveiling the Perils of the TorchScript Engine in PyTorch

Ji'an ZhouLishuo Song

Discloses multiple memory corruption and deserialization vulnerabilities in PyTorch's TorchScript JIT engine that allow arbitrary remote code execution via malicious AI models.

Summary

PyTorch TorchScript and JIT Execution Engine

PyTorch TorchScript allows machine learning models to be serialized into standalone C++ runtimes for high-performance production inference across cloud servers and edge devices.

Analyzing the TorchScript Intermediate Representation

Ji'an Zhou and Lishuo Song audit the TorchScript JIT compiler, bytecode interpreter, and model deserialization routines, uncovering critical memory corruption and type confusion flaws.

Remote Code Execution via Weaponized Models

Crafting malicious model files targeting TorchScript parser vulnerabilities allows attackers to execute arbitrary code on cloud inference servers upon loading the model, bypassing Python sandbox restrictions.

Securing Machine Learning Model Pipelines

Remediating TorchScript vulnerabilities requires ML developers to treat untrusted model weights as untrusted binaries, enforce strict deserialization schemas, and run inference within isolated containers.

Related talks