PT2 存档规范

创建日期：2025年7月16日 | 最后更新日期：2025年12月3日

以下规范定义了可通过以下方法生成的存档格式：

• torch.export，通过调用 torch.export.save()

• AOTInductor，通过调用 torch._inductor.aoti_compile_and_package()

该存档是一个 zip 文件，可以使用标准的 zipfile API 进行操作。

以下是一个存档示例。我们将逐个文件夹对存档内容进行解析。

.
├── archive_format
├── byteorder
├── .data
│   ├── serialization_id
│   └── version
├── data
│   ├── aotinductor
│   │   └── model1
│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.kernel_metadata.json
│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.kernel.cpp
│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.wrapper_metadata.json
│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.wrapper.cpp
│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.wrapper.so
│   │       ├── cg7domx3woam3nnliwud7yvtcencqctxkvvcafuriladwxw4nfiv.cubin
│   │       └── cubaaxppb6xmuqdm4bej55h2pftbce3bjyyvljxbtdfuolmv45ex.cubin
│   ├── weights
│   │  ├── model1_weights_config.json
│   │  ├── model2_weights_config.json
│   │  ├── weight_0
│   │  ├── weight_1
│   │  ├── weight_2
│   └── constants
│   │  ├── model1_constants_config.json
│   │  ├── model2_constants_config.json
│   │  ├── tensor_0
│   │  ├── tensor_1
│   │  ├── custom_obj_0
│   │  ├── custom_obj_1
│   └── sample_inputs
│       ├── model1.pt
│       └── model2.pt
├── extra
│   └── ....json
└── models
    ├── model1.json
    └── model2.json

内容

存档头信息

• archive_format 声明了此存档使用的格式。目前仅支持 “pt2”。

• byteorder。为 “little” 或 “big” 中的一个，供 zip 文件读取器使用。

• /.data/version 包含存档版本。（请注意，这既不是导出序列化的模式版本，也不是 Aten Opset 版本）。

• /.data/serialization_id 是为当前存档生成的哈希值，用于验证。

AOTInductor 编译制品

路径：/data/aotinductor/<model_name>-<backend>/

AOTInductor 编译制品针对每个“模型-后端”对进行保存。例如，模型 model1 在 A100 和 H100 上的编译制品将分别保存在 model1-a100 和 model1-h100 文件夹中。

该文件夹通常包含：

• <uuid>.wrapper.so：由 .cpp 编译生成的动态链接库。.cpp。

• <uuid>.wrapper.cpp：AOTInductor 生成的 C++ 封装文件。

• <uuid>.kernel.cpp：AOTInductor 生成的 C++ 内核文件。

• *.cubin：由 Triton 代码生成内核编译而成的 Triton 内核。

• <uuid>.wrapper_metadata.json：从 aot_inductor.metadata Inductor 配置中传递的元数据。

• (可选) <uuid>.json：由 ProxyExecutor 执行的自定义算子的外部回退节点，根据 ExternKernelNode 结构进行序列化。如果模型未使用自定义算子/ProxyExecutor，则省略此文件。

权重

路径：/data/weights/*

模型参数和缓冲区保存在 /data/weights/ 文件夹中。每个张量被保存为一个独立的文件。该文件仅包含原始数据块，张量元数据以及从模型权重 FQN 到所保存原始数据块的映射则单独保存在 <model_name>_weights_config.json 中。

常量

路径：/data/constants/*

TensorConstants（张量常量）、非持久化缓冲区和 TorchBind 对象保存在 /data/constants/ 文件夹中。元数据以及从模型常量 FQN 到所保存原始数据块的映射则单独保存在 <model_name>_constants_config.json 中。

示例输入

路径：/data/sample_inputs/<model_name>.pt

torch.export 使用的 sample_input 可以包含在存档中以供后续使用。通常，它是一个张量扁平化列表，结合了 forward() 函数的 args 和 kwargs。

该 .pt 文件由 torch.save(sample_input) 生成，并可通过 Python 中的 torch.load() 和 C++ 中的 torch::pickle_load() 加载。

当模型拥有多个示例输入副本时，它们将被打包为 <model_name>_<index>.pt。

模型定义

路径：/models/<model_name>.json

模型定义是来自 torch.export.save 的 ExportedProgram 的序列化 JSON 以及其他模型级元数据。

多模型支持

此存档规范支持在同一个文件中并存多个模型定义，其中 <model_name> 作为模型的唯一标识符，并将作为存档其他文件夹中的引用使用。

更底层的 API（如 torch.export.pt2_archive._package.package_pt2() 和 torch.export.pt2_archive._package.load_pt2()）允许您对打包和加载过程进行更细粒度的控制。