ivanphuc/pytorch-mmap-weightsonly-poc

PyTorch mmap weights_only security PoC

Overview

This repository contains a malformed PyTorch ZIP checkpoint demonstrating that the same archive is:

• rejected by torch.load(..., weights_only=True, mmap=False)
• accepted by torch.load(..., weights_only=True, mmap=True)

File

• offset_patched.pt: malformed checkpoint proof of concept

Tested environment

• Python 3.12
• PyTorch 2.12.1+cpu

Reproduction

import torch

torch.load("offset_patched.pt", weights_only=True, mmap=False)
# expected: RuntimeError due to record size mismatch

obj = torch.load("offset_patched.pt", weights_only=True, mmap=True)
print(obj["base"].tolist())
print(obj["view"].tolist())
print(obj["base"].untyped_storage().nbytes())

Expected result

• mmap=False rejects the malformed checkpoint
• mmap=True accepts it
• loaded tensors contain attacker-controlled / malformed values
• loaded storage size reflects forged metadata rather than the legitimate record payload size