Open-Source Wikis

/

PyTorch

/

Reference

/

Data models

pytorch/pytorch

Data models

The serialization formats and IRs PyTorch reads and writes.

.pt / .pth (the eager checkpoint)

Container produced by torch.save:

mymodel.pt   ← actually a zip file
├── data.pkl                 pickled object graph (with reduced Tensors)
├── data/
│   ├── 0                    raw tensor bytes for storage 0
│   ├── 1                    raw tensor bytes for storage 1
│   └── ...
├── version                  serialization protocol version
└── byteorder                little/big endian indicator

Properties:

  • The pickle stream uses reducer hooks (torch._tensor._tensor_reduce) so a Tensor becomes (rebuild_tensor, (storage_offset, sizes, strides, requires_grad, ...)).
  • Storages live in data/N files, one per unique storage. The pickle references storages by index. Multiple tensors sharing storage share the file.
  • Loading with weights_only=True uses a restricted unpickler that disallows arbitrary globals.

TorchScript (.pt / .ptl)

Same zip container, but with additional entries:

model.pt
├── code/                    serialized Python source for each scripted method
├── constants.pkl            constants referenced by code
├── data.pkl                 module structure
├── data/                    storages
├── extra/                   user-supplied extra files
└── version, byteorder

Loaded with torch::jit::load (C++) or torch.jit.load (Python) — produces a torch::jit::Module with executable methods.

ONNX

.onnx files: protobuf encoding of the ONNX graph. PyTorch produces them via torch.onnx.export (legacy tracer-based) or torch.onnx.dynamo_export (Dynamo + ONNX-script). The exported model is portable to any ONNX runtime.

ExportedProgram (torch.export)

In-memory IR returned by torch.export.export(model, args):

ExportedProgram
├── graph_module: torch.fx.GraphModule    nodes are aten ops, ATen IR
├── graph_signature: ExportGraphSignature  describes inputs/outputs/buffers/parameters
├── state_dict: dict[str, Tensor]
├── range_constraints: dict[Symbol, ValueRanges]
├── module_call_graph: list[ModuleCallEntry]
└── verifier: Verifier                      ATen IR validity checker

Serialized to .pt2 files via torch.export.save — which produces another zip container with the FX graph as a JSON-ish protobuf and the state dict as raw tensors.

AOTInductor .so package

Output of torch._inductor.aot_compile(exported_program, ...):

some_model.pt.package/
├── model.so                C++ generated kernel + dispatch graph
├── model.json              serialized kernel metadata
└── ...

Loaded from C++ via torch::inductor::AOTIModelContainerRunner. No Python required.

Distributed checkpoint (DCP)

torch.distributed.checkpoint writes per-rank metadata + sharded payload to a directory:

checkpoint/
├── .metadata               Plan and shape descriptions
├── __0_0.distcp            payload chunk for rank 0
├── __1_0.distcp            payload chunk for rank 1
└── ...

DCP supports resharding — you can load a checkpoint saved on N ranks onto M ranks. Useful for elastic training.

ATen IR

Not a file format; an invariant the export and compile stacks observe. After make_fx, functionalization, and decomposition, an FX graph contains only:

  • aten.* op calls (the core_aten opset, ~250 ops).
  • Function I/O (placeholder, output).
  • Get-attr nodes for parameters / buffers.

This is the canonical IR for downstream backends.

Pre-dispatch IR

A wider IR used by export's "pre-dispatch" mode: includes higher-level functional ops (call_function to torch.nn.functional.linear, torch.ops.higher_order.cond, etc.) before they're decomposed to ATen.

Triton kernels

Inductor-emitted kernels are persisted to disk as .py files (the Triton source) plus a .cubin/.json (the compiled binary + launch metadata). The cache directory defaults to ~/.cache/torch/inductor/ and can be redirected via TORCHINDUCTOR_CACHE_DIR.

Fx GraphModule

Not a file format but worth noting: torch.fx.GraphModule is the in-memory IR shared by Dynamo, AOTAutograd, Inductor, ONNX export, quantization, FX-based passes. It's a Graph (list of Nodes, each with op/target/args/kwargs) bundled with the parameter/buffer state and a generated forward Python.

Where to look

Format File
.pt / .pth torch/serialization.py, torch/_weights_only_unpickler.py
TorchScript torch/csrc/jit/serialization/
ONNX torch/onnx/
ExportedProgram torch/export/
AOTInductor torch/_inductor/aot_inductor.py, torch/csrc/inductor/aoti_runtime/
DCP torch/distributed/checkpoint/
FX torch/fx/

Built by Factory AutoWiki from public repository content. It is a generated preview for codebase exploration, not source-maintained documentation.

Data models – PyTorch wiki | Factory