PT2 归档规范

创建日期：2025年7月16日 | 最后更新日期：2025年9月09日

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

• 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/<模型名称>-<后端>/

AOTInductor 编译产物为每个模型-后端对保存。例如，A100和H100上model1模型的编译产物将分别保存在model1-a100和model1-h100文件夹中。

文件夹通常包含

• <uuid>.wrapper.so：从 .cpp 编译的动态库。.cpp。

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

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

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

• <uuid>.wrapper_metadata.json：从 aot_inductor.metadata 归纳器配置传入的元数据。

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

权重

路径：/data/weights/*

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

常量

路径：/data/constants/*

TensorConstants，非持久性缓冲区和TorchBind对象保存在 /data/constants/ 文件夹中。元数据和从模型常量 FQN 到已保存原始数据块的映射单独保存在 <模型名称>_constants_config.json 中。

样本输入

路径：/data/sample_inputs/<模型名称>.pt

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

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

当模型有多个样本输入副本时，它将被打包为 <模型名称>_<索引>.pt。

模型定义

路径：/models/<模型名称>.json

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

多个模型

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

像 torch.export.pt2_archive._package.package_pt2() 和 torch.export.pt2_archive._package.load_pt2() 这样的低级 API 允许您更精细地控制打包和加载过程。