yibai777/pytorch-audit2

PyTorch PT2 Archive RCE via weights_only=False (Post-fix Bypass)

Format: Torch Export (.pt2) – PyTorch
Project: PyTorch (pytorch/pytorch)
Version: 2.12.0
Severity: Critical – RCE on model load

---

CVE Context

• CVE-2024-31580 (PyTorch < 2.2.2) covered the general torch.export.load deserialization issue. The fix added a restricted pickle allowlist for the exported program metadata.
• This finding is a post-fix bypass: the WEIGHTS loading path at _package.py:878 uses torch.load(weights_only=False) directly — a separate code path not covered by the CVE-2024-31580 fix.
• The vulnerability is still present in PyTorch 2.12.0 (latest stable).

Description

torch.export.load() loads model weights from within a PT2 archive using torch.load() with weights_only=False:

torch/export/pt2_archive/_package.py:878:

state_dict[weight_fqn] = torch.load(
    io.BytesIO(weight_bytes), weights_only=False
)

Even though torch.load() defaults to weights_only=True since PyTorch 2.6, the PT2 archive loader explicitly passes weights_only=False, allowing arbitrary code execution when loading a malicious .pt2 file.

Impact

Any user or application that calls torch.export.load() on an untrusted .pt2 file gets RCE. This includes users downloading PT2 models from Hugging Face, model zoos, or any untrusted source.

Steps to Reproduce

# 1. Generate malicious .pt2 (executes calc.exe by default)
python poc_pt2_rce.py --output malicious.pt2

# 2. Victim loads the file
torch.export.load("malicious.pt2")

Or test directly:

python poc_pt2_rce.py --cmd "echo PWNED" --test-load

How It Works

1. Export a legitimate model → PT2 archive
2. Modify weights_config to set use_pickle=True
3. Replace weight data with a PyTorch zip containing a malicious pickle (os.system)
4. When loaded, _load_state_dict calls torch.load(weights_only=False) → RCE

Files

File	Purpose
poc_pt2_rce.py	Exploit PoC – generate & test malicious .pt2
malicious.pt2	Pre-built malicious archive (executes calc.exe)

Reference

• torch/export/pt2_archive/_package.py:878
• torch/serialization.py:1316-1637