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 允許您更精細地控制打包和載入過程。