Serialization

Active contributors: mikaylagawarecki

Purpose

Saving and loading models and tensors is a deceptively complicated topic in PyTorch. Three formats coexist:

1. Python pickle (torch.save/torch.load) — the default, most flexible, and (until recently) the most dangerous because pickle can execute arbitrary code on load.
2. Weights-only loading — a restricted unpickler that only allows torch.Tensor and other vetted types. This is now the default in PyTorch 2.x for torch.load.
3. TorchScript / ZIP archive (torch.jit.save) — a structured zip with serialized IR + weights. The format used by PyTorch Mobile and most pre-2.0 production deployments.

A fourth, newer path is torch.export + torch._export.save/load, which serializes an ExportedProgram for AOTInductor and ExecuTorch.

Directory layout

Path	Contents
torch/serialization.py	torch.save, torch.load, the user-facing API (~87K lines)
torch/storage.py	Storage Python wrapper
torch/_weights_only_unpickler.py	Restricted unpickler
torch/csrc/Storage*.cpp	C++ Storage Python binding
torch/csrc/jit/serialization/	TorchScript zip format
caffe2/serialize/	Low-level zip / file reader/writer (PyTorchFileReader/Writer)
torch/_export/serde/	torch.export schema (Pydantic) and serialization
aten/src/ATen/MapAllocator.cpp	mmap-backed storage for shared/large tensors

Key abstractions

Type	File	Purpose
torch.save / torch.load	torch/serialization.py	Top-level API
_open_zipfile_writer	torch/serialization.py	The new (PyTorch 1.6+) zip-based format
_open_legacy_zip	torch/serialization.py	The legacy tar-based format
Unpickler (weights-only)	torch/_weights_only_unpickler.py	Restricted Python unpickler
PyTorchFileReader/Writer	caffe2/serialize/inline_container.h	Low-level zip stream reader/writer
_load_jit_data	torch/jit/_serialization.py	TorchScript loader

How it works

torch.save / torch.load (modern)

The default save format is a zip archive containing:

• data.pkl — pickled top-level object (e.g., a state_dict or a model). Pickle uses persistent IDs for tensors, so each tensor in the pickle stream is replaced with a small marker.
• data/<n> — one file per unique storage, containing the raw bytes.
• byteorder — little-endian indicator.
• (For torch.jit.save only) code/, constants.pkl, etc.

Saving streams the pickled object first, then writes each storage referenced by a persistent ID into a separate zip entry. Loading is the inverse: pickle reads markers, the loader maps each marker to a torch.Storage view of the corresponding zip entry, and tensors are reconstructed with the right shape/strides.

graph LR
    User[Python object<br/>state_dict / model] -->|pickle.dump<br/>w/ persistent_id| Pkl[data.pkl]
    User -->|raw bytes per storage| Zip[data/0, data/1, ...]
    Pkl --> Archive[ZIP archive]
    Zip --> Archive

Weights-only

torch.load(..., weights_only=True) (the default since 2.x for tensors) uses torch/_weights_only_unpickler.py instead of stock pickle.Unpickler. The restricted unpickler:

• Allows only a hard-coded set of types: torch.Tensor, torch.Storage, torch.dtype, torch.device, collections.OrderedDict, simple types.
• Disallows arbitrary GLOBAL opcodes; users can opt-in via torch.serialization.add_safe_globals.
• Refuses REDUCE calls except for the vetted constructors.

This closes the long-standing remote-code-execution gotcha where anyone with a .pt file could ship a malicious payload.

torch.jit.save

For ScriptModules the format is the same zip archive but with extra entries:

• code/<class>.py — TorchScript source for the model.
• constants.pkl — pickled constants.
• attributes.pkl — module attributes.
• version — serialization format version.

Loaded via torch.jit.load, which goes through torch/csrc/jit/serialization/import.cpp and reconstructs the IR.

torch.export save/load

torch._export.save(ep, path) writes a directory or zip with the FX graph (in textual form), the schema-validated metadata (torch/_export/serde/schema.py), and parameter/buffer tensors. The format is versioned via the schema; backward compatibility is enforced by tests in test/export/.

Maps and mmap

aten/src/ATen/MapAllocator.cpp implements memory-mapped tensors. torch.load(..., mmap=True) mmap-loads each storage instead of reading it into RAM, which is essential for very large LLM checkpoints.

Integration points

• Distributed. Distributed checkpoint (torch.distributed.checkpoint) bypasses torch.save for large jobs and uses its own per-shard format. See Distributed.
• Mobile. PyTorch Mobile reads .ptl (lite-interpreter) files, an extension of the JIT zip format with a stripped runtime.
• HuggingFace safetensors is preferred by some communities for tensor-only data; PyTorch interoperates via the safetensors package, not native code.
• Security. Per SECURITY.md, weights-only load is the supported safe path; the old default (arbitrary pickle) should be considered untrusted.

Entry points for modification

• New tensor type to support in weights-only load → register via torch.serialization.add_safe_globals (user) or add to the vetted list in _weights_only_unpickler.py (in-tree).
• Format changes → bump the _pyzipfile_data_pkl_format_version in torch/serialization.py and update readers.
• New mmap behaviour → aten/src/ATen/MapAllocator.cpp.

Key source files

File	Purpose
torch/serialization.py	User-facing save/load
torch/_weights_only_unpickler.py	Restricted unpickler
caffe2/serialize/inline_container.cc	Zip reader/writer
torch/csrc/jit/serialization/pickler.cpp	TorchScript pickler
torch/csrc/jit/serialization/import.cpp	TorchScript loader
aten/src/ATen/MapAllocator.cpp	mmap-backed storage
torch/_export/serde/schema.py	torch.export serialization schema