PyTorch PT2 Archive 规范深度解析torch.export 与 AOTInductor 的统一模型打包格式【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文以 PyTorch 官方文档 PT2 Archive Spec 为骨架结合仓库内 打包实现、目录常量定义 与 高层 API 源码系统讲解 PT2 Archive 的目录结构、各组成部分的存储约定、多模型支持机制以及torch.export.save/torch.export.load/package_pt2/load_pt2的完整用法。读完本文你将能够理解.pt2文件内部每个文件夹的含义并具备自行打包、解包、校验与加载 PT2 Archive 的实战能力。一、PT2 Archive 是什么PT2 Archive 是 PyTorch 2 系列定义的一种统一归档格式archive format用于把导出模型及其配套产物打包进单个文件方便分发、部署与下游消费。按官方规范该归档可以通过以下两种方式产生通过 {func}torch.export的torch.export.save保存ExportedProgram通过torch._inductor.aoti_compile_and_package打包 AOTInductor 编译产物。归档文件本身就是一个 zipfile实际由torch._C.PyTorchFileWriter写出因此可以使用标准的 zipfile API 直接查看、解压与操作内部条目。在源码层面判断一个文件是否是 PT2 Archive 由is_pt2_package完成见 torch/export/pt2_archive/_package.py#L74-L90它用zipfile.ZipFile打开文件读取根目录下的archive_format记录判断其内容是否为bpt2。这也是 PT2 Archive 的身份证——格式标识本身是归档头信息的一部分。二、完整目录结构总览规范给出了一个包含两个模型model1、model2的示例归档原文目录树如下. ├── 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把目录与源码常量对应起来常量在 torch/csrc/export/pt2_archive_constants.h#L8-L62 中集中定义再经 constants.py 暴露给 Python 层目录/文件源码常量含义archive_formatARCHIVE_FORMAT_PATH/ARCHIVE_FORMAT_VALUE格式标识固定为pt2archive_versionARCHIVE_VERSION_PATH/ARCHIVE_VERSION_VALUE归档规范版本号初始为0models/{}.jsonMODELS_FILENAME_FORMAT模型定义ExportedProgram 序列化data/aotinductor/AOTINDUCTOR_DIRAOTInductor 编译产物data/weights/WEIGHTS_DIR权重参数 buffer原始数据data/weights/{}_weights_config.jsonWEIGHTS_CONFIG_FILENAME_FORMAT权重 FQN → 数据文件映射data/constants/CONSTANTS_DIRTensor 常量、非持久 buffer、TorchBind 对象data/constants/{}_constants_config.jsonCONSTANTS_CONFIG_FILENAME_FORMAT常量 FQN → 数据文件映射data/constants/tensor_*TENSOR_CONSTANT_FILENAME_PREFIXTensor 常量数据文件前缀data/constants/custom_obj_*CUSTOM_OBJ_FILENAME_PREFIXTorchBindScriptObject数据文件前缀data/constants/opaque_obj_*OPAQUE_OBJ_FILENAME_PREFIX不透明自定义对象数据文件前缀data/sample_inputs/{}.ptSAMPLE_INPUTS_FILENAME_FORMAT样例输入data/executorch/EXECUTORCH_DIRExecuTorch 产物如 PTE 文件extra/EXTRA_DIR用户自定义附加文件data/mtiaMTIA_DIRMTIA 加速器产物需要说明的是archive_version的记录位于归档根目录常量值为archive_version而规范中.data/version一节所描述的是 zip 文件格式自身的版本源码注释明确区分了两者——PT2ArchiveWriter调用set_min_version(6)写入的是 zip 格式版本写向/.data/versionPT2 Archive 规范版本则写向/archive_version见 torch/export/pt2_archive/_package.py#L93-L117。三、Archive 头信息Archive Headersarchive_format声明归档使用的格式。目前只允许取值pt2。写入与校验都在源码中硬编码PT2ArchiveWriter.__exit__在关闭归档前若未写入该记录会自动补写ARCHIVE_FORMAT_VALUE而PT2ArchiveReader构造时会校验read_string(ARCHIVE_FORMAT_PATH) ! ARCHIVE_FORMAT_VALUE并抛出AssertionError(Invalid archive format)见 torch/export/pt2_archive/_package.py#L195-L211。byteorder取值为little或big供 zip 文件读取端判断字节序。/.data/version与/archive_version规范中/.data/version包含归档版本号并特别强调它既不是 export 序列化的 schema 版本也不是 ATen Opset 版本。从实现看PT2 Archive 规范自身的版本号记录在根目录archive_version文件中初始值为02024 年 9 月 4 日确立的初始版本。load_pt2加载时会读取archive_version并与当前ARCHIVE_VERSION_VALUE比对不一致即抛出ValueError见 torch/export/pt2_archive/_package.py#L1136-L1141。源码注释对版本号升级规则有明确约定当 PT2 Archive 规范出现破坏性变更例如删除某个文件夹、改变下列字段的命名约定时必须递增ARCHIVE_VERSION_VALUE且 Archive 读取端需要相应支持同时加载当前与旧版本的 PT2 Archive见 torch/csrc/export/pt2_archive_constants.h#L17-L24。/.data/serialization_id为当前归档生成的哈希用于校验。归档读取端PT2ArchiveReader同样暴露了archive_version()方法——当archive_version记录缺失时会将其视为版本 0从而保持对早期归档的向后兼容见 torch/export/pt2_archive/_package.py#L229-L240。四、AOTInductor 编译产物data/aotinductor/路径约定路径格式为/data/aotinductor/model_name-backend/。AOTInductor 编译产物按模型-后端组合分别保存例如model1模型在 A100 与 H100 上的编译产物会分别存放到model1-a100与model1-h100文件夹中互不干扰。文件夹典型内容该文件夹通常包含以下文件uuid.wrapper.so由uuid.cpp编译生成的动态库uuid.wrapper.cppAOTInductor 生成的 C wrapper 文件uuid.kernel.cppAOTInductor 生成的 C kernel 文件*.cubin由 Triton codegen kernel 编译出的 Triton kernel 二进制uuid.wrapper_metadata.json从aot_inductor.metadatainductor 配置传入的元数据可选uuid.json自定义算子需要由ProxyExecutor执行时的外部 fallback 节点按ExternKernelNode结构序列化。若模型未使用自定义算子/ProxyExecutor则该文件省略。打包实现细节在 torch/export/pt2_archive/_package.py#L254-L328 的_package_aoti_files中可以看到更精细的逻辑aoti_files可以传列表自动命名为model或{模型名: 文件列表}字典文件按 basename 判断去向以custom_obj_前缀开头的文件进入data/constants/其余文件进入data/aotinductor/model_name/源码注释特别指出早期实现拒绝一个模型存在多个.so的情况但Triton CPU AOTInductor 每个模型会产生多个.so文件wrapper、launcher、kernel因此当前实现支持多.so若存在权重会对权重做去重分组group_weights将共享同一 storage 的多个张量合并保存并在data/aotinductor/model_name/weights_config.json中记录weight_name - (filename, shape, stride, offset)的映射加载时通过as_strided还原见 torch/export/pt2_archive/_package.py#L1208-L1222。五、权重存储data/weights/存储约定路径为/data/weights/*。模型的参数parameters与缓冲区buffers保存在该文件夹中每个张量保存为独立文件文件内只有原始数据 blobraw data blob。张量的元数据以及模型权重 FQN → 原始数据 blob的映射单独保存在model_name_weights_config.json中。打包端两种序列化路径_package_state_dicttorch/export/pt2_archive/_package.py#L435-L479会遍历exported_program.state_dict把权重分成两类pickle 路径张量是 tensor subclass 且非 fake tensor 时_should_use_pickle见 L366-L367用torch.save保存整个张量对象文件名如weight_idxraw bytes 路径普通张量通过_get_raw_tensor_bytes直接拷贝untyped_storage的原始字节见 L339-L363并对共享 storage 的张量做去重合并。每条权重在 config 中记录为schema.PayloadMeta包含path_name、is_param是否是torch.nn.Parameter、use_pickle标志以及序列化的tensor_metashape/stride/storage_offset/dtype/device 等。加载端按 FQN 还原并保留别名关系_load_state_dicttorch/export/pt2_archive/_package.py#L864-L917先用_build_file_map把每个weight_*文件读成 flat 张量再对每个权重 FQN 用torch.as_strided(size, stride, storage_offset)从 flat 张量上切出正确的视图is_param为真的张量恢复为torch.nn.Parameter并还原requires_grad。源码注释指出这种FQN → 文件 → strided 张量的链式映射保留了权重之间的别名aliasing关系。另外加载端对旧格式做了兼容若归档中存在data/weights/model_name.pt这一历史遗留文件会走 legacy 分支读取并给出警告提示使用torch.export.save重新生成新格式见 L868-L875。六、常量存储data/constants/存储约定路径为/data/constants/*。三类对象保存在此Tensor 常量TensorConstants非持久缓冲区non-persistent buffersTorchBind 对象。元数据与模型常量 FQN → 原始数据 blob的映射保存在model_name_constants_config.json中。打包端三类前缀与两种格式_package_constantstorch/export/pt2_archive/_package.py#L482-L565按对象类型分类处理Tensor 常量torch.Tensor与权重类似subclass 走 pickle 路径普通张量走 raw bytes 路径文件名前缀为tensor_ScriptObjecttorch._C.ScriptObject使用 torchbind 感知的 picklerzip-archive 格式加载端可在 C 中读取文件名前缀为custom_obj_不透明自定义对象custom class object使用标准 Python pickle文件名前缀为opaque_obj_。源码注释特别说明不同的文件名前缀正是加载端区分格式的判别标志discriminator。加载端_load_constantstorch/export/pt2_archive/_package.py#L920-L985依据path_name的前缀分发tensor_前缀走与权重一致的 as_strided 还原或torch.loadcustom_obj_前缀用torch._C._pickle_load_obj反序列化opaque_obj_前缀用标准pickle.loads。若遇到未知前缀则抛出RuntimeError。同样地若归档中存在data/constants/model_name.pt旧格式文件也会走 legacy 兼容分支。七、样例输入data/sample_inputs/路径为/data/sample_inputs/model_name.pt。torch.export使用的样例输入sample input可以一并打包进归档供下游使用。通常它是一个扁平化的张量列表同时包含forward()的 args 与 kwargs实际打包时通过reorder_kwargs保证 kwargs 顺序与调用约定一致见 torch/export/pt2_archive/_package.py#L33。.pt文件由torch.save(sample_input)产生读取方式为Python 端torch.load()C 端torch::pickle_load()。当模型存在多份样例输入时会打包为model_name_index.pt如model1_0.pt、model1_1.pt。八、模型定义models/路径为/models/model_name.json。模型定义是torch.export.save产出的ExportedProgram 序列化 JSON外加其他模型级元数据。从打包流程看torch/export/pt2_archive/_package.py#L581-L626每个模型会依次写入data/weights/model_name_weights_config.json权重 payload 配置data/constants/model_name_constants_config.json常量 payload 配置models/model_name.jsonserialize()产出的 ExportedProgram 序列化数据此时serialize_state_dictFalse、serialize_constantsFalse即权重/常量不重复内嵌data/sample_inputs/model_name.pt样例输入。加载端_load_exported_programsL988-L1026扫描所有models/前缀条目从文件名反推模型名读取样例输入、反序列化 ExportedProgram再把上一步还原的state_dict与constants一并交给ExportedProgramDeserializer完成整体重建。九、多模型共存PT2 Archive 规范支持多个模型定义共存于同一个文件中。model_name是模型的唯一标识符并被归档其他文件夹作为引用依据——例如data/weights/model1_weights_config.json对应models/model1.jsondata/aotinductor/model1/对应model1的编译产物。在 Python API 层面传入dict[str, ExportedProgram]或dict[str, list[str]]AOTI 文件即可打包多个模型只传单个对象时自动命名为model见 torch/export/pt2_archive/_package.py#L590-L596。加载端返回的PT2ArchiveContents中exported_programs与aoti_runners均为以模型名为键的字典L788-L792。十、高层与底层 APIsave/load 与 package_pt2/load_pt2高层 APItorch.export.save与torch.export.loadtorch.export.savetorch/export/init.py#L266-L335内部就是对package_pt2的封装官方 docstring 给出了完整可运行示例import torch import io class MyModule(torch.nn.Module): def forward(self, x): return x 10 ep torch.export.export(MyModule(), (torch.randn(5),)) # 保存到文件 torch.export.save(ep, exported_program.pt2) # 保存到 io.BytesIO 缓冲区 buffer io.BytesIO() torch.export.save(ep, buffer) # 携带附加文件保存 extra_files {foo.txt: bbar.decode(utf-8)} torch.export.save(ep, exported_program.pt2, extra_filesextra_files)对应的torch.export.loadtorch/export/init.py#L338-L405import torch import io # 从文件加载 ExportedProgram ep torch.export.load(exported_program.pt2) # 从 io.BytesIO 对象加载 with open(exported_program.pt2, rb) as f: buffer io.BytesIO(f.read()) buffer.seek(0) ep torch.export.load(buffer) # 加载附加文件map 中的 value 会被替换为实际数据 extra_files {foo.txt: } ep torch.export.load(exported_program.pt2, extra_filesextra_files) print(extra_files[foo.txt]) print(ep(torch.randn(5)))需要注意两个官方 warning该 API 处于活跃开发阶段旧版本保存的文件可能无法在新版本中使用且torch.export.load底层使用 pickle切勿加载不可信来源的数据。底层 APIpackage_pt2与load_pt2package_pt2torch/export/pt2_archive/_package.py#L648-L728提供更细粒度的打包控制支持以下参数参数类型说明fstr / os.PathLike / IO[bytes]输出文件或可写缓冲区推荐以.pt2结尾exported_programsExportedProgram或dict[str, ExportedProgram]要保存的导出程序可多模型aoti_fileslist[str]或dict[str, list[str]]torch._inductor.aot_compile(..., {aot_inductor.package: True})生成的 AOTInductor 文件extra_filesdict[str, Any]附加文件写入extra/目录opset_versiondict[str, int]opset 名到版本号的映射pickle_protocolint覆盖默认 pickle 协议默认值为 2executorch_filesdict[str, bytes]可选的 ExecuTorch 产物若三者都未传入exported_programs、aoti_files、extra_files全为None会直接抛出断言错误。打包全程在PT2ArchiveWriter上下文管理器中完成关闭时自动补写archive_format与archive_version头记录。load_pt2torch/export/pt2_archive/_package.py#L1082-L1224的加载参数参数默认值说明f—文件路径或可读缓冲区expected_opset_versionNone期望的 opset 版本映射num_runners1加载 AOTInductor 产物的 runner 数量run_single_threadedFalse是否以单线程模式运行避免与 CUDA Graph 的同步逻辑冲突device_index-1加载目标设备索引CUDA 下-1对应cuda传1则加载到cuda:1load_weights_from_diskFalse是否把权重解出到磁盘再由 AOTI runner 加载加载后得到PT2ArchiveContentsexported_programs、aoti_runners、extra_files。其中aoti_runners里的每个元素是AOTICompiledModelL731-L785它是一个可调用对象调用时会取出AOTIModelPackageLoader的 call spec用pytree对输入输出做展平/还原然后执行boxed_run还提供了get_metadata()、get_constant_fqns()以及load_constants(constants_map, check_full_update..., user_managed..., allow_h2d_copy...)用于运行期更新常量。user_managedTrue时 loader 直接保存张量指针调用方必须保证其存活allow_h2d_copyTrue时 CPU 张量会被静默拷贝到模型所在设备便于把 CPU 的state_dict()加载进非 CPU 模型与user_managed不兼容。需要说明的是当传入的是内存缓冲区且归档内含 AOTI 产物时load_pt2会先把缓冲区内容写入临时.pt2文件再交给AOTIModelPackageLoader因后者不支持直接读缓冲区见 L1175-L1195。十一、版本与兼容性归档规范版本当前ARCHIVE_VERSION_VALUE为0后续任何破坏性变更删目录、改命名约定等都必须递增该值且读取端需要同时兼容新旧版本见 torch/csrc/export/pt2_archive_constants.h#L15-L24。load_pt2在版本不匹配时直接报ValueError。zip 格式版本PT2ArchiveWriter将 zip 最小版本设为 6该版本写入/.data/version与归档规范版本是两套独立体系。遗留格式兼容data/weights/model_name.pt与data/constants/model_name.pt属于历史遗留格式加载时会打印 warning 并提示改用torch.export.save重新生成新格式。读取端校验PT2ArchiveReader构造时校验archive_formatis_pt2_package可用来无损检测任意字节流/路径是否为 PT2 Archive适合在部署管线中做格式分发决策。总体而言PT2 Archive 通过头信息 按模型组织的分区目录 JSON 配置索引原始数据 blob的设计把ExportedProgram的图定义、样例输入、AOTInductor 编译产物、权重与常量统一收纳进一个 zip 文件既保持了标准 zip 的可操作性又通过model_name约定天然支持多模型、多后端产物共存是理解 PyTorch 2 系列导出与部署链路的关键规范。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyTorch PT2 Archive 规范深度解析torch.export 与 AOTInductor 的统一模型打包格式【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文以 PyTorch 官方文档 PT2 Archive Spec 为骨架结合仓库内 打包实现、目录常量定义 与 高层 API 源码系统讲解 PT2 Archive 的目录结构、各组成部分的存储约定、多模型支持机制以及torch.export.save/torch.export.load/package_pt2/load_pt2的完整用法。读完本文你将能够理解.pt2文件内部每个文件夹的含义并具备自行打包、解包、校验与加载 PT2 Archive 的实战能力。一、PT2 Archive 是什么PT2 Archive 是 PyTorch 2 系列定义的一种统一归档格式archive format用于把导出模型及其配套产物打包进单个文件方便分发、部署与下游消费。按官方规范该归档可以通过以下两种方式产生通过 {func}torch.export的torch.export.save保存ExportedProgram通过torch._inductor.aoti_compile_and_package打包 AOTInductor 编译产物。归档文件本身就是一个 zipfile实际由torch._C.PyTorchFileWriter写出因此可以使用标准的 zipfile API 直接查看、解压与操作内部条目。在源码层面判断一个文件是否是 PT2 Archive 由is_pt2_package完成见 torch/export/pt2_archive/_package.py#L74-L90它用zipfile.ZipFile打开文件读取根目录下的archive_format记录判断其内容是否为bpt2。这也是 PT2 Archive 的身份证——格式标识本身是归档头信息的一部分。二、完整目录结构总览规范给出了一个包含两个模型model1、model2的示例归档原文目录树如下. ├── 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把目录与源码常量对应起来常量在 torch/csrc/export/pt2_archive_constants.h#L8-L62 中集中定义再经 constants.py 暴露给 Python 层目录/文件源码常量含义archive_formatARCHIVE_FORMAT_PATH/ARCHIVE_FORMAT_VALUE格式标识固定为pt2archive_versionARCHIVE_VERSION_PATH/ARCHIVE_VERSION_VALUE归档规范版本号初始为0models/{}.jsonMODELS_FILENAME_FORMAT模型定义ExportedProgram 序列化data/aotinductor/AOTINDUCTOR_DIRAOTInductor 编译产物data/weights/WEIGHTS_DIR权重参数 buffer原始数据data/weights/{}_weights_config.jsonWEIGHTS_CONFIG_FILENAME_FORMAT权重 FQN → 数据文件映射data/constants/CONSTANTS_DIRTensor 常量、非持久 buffer、TorchBind 对象data/constants/{}_constants_config.jsonCONSTANTS_CONFIG_FILENAME_FORMAT常量 FQN → 数据文件映射data/constants/tensor_*TENSOR_CONSTANT_FILENAME_PREFIXTensor 常量数据文件前缀data/constants/custom_obj_*CUSTOM_OBJ_FILENAME_PREFIXTorchBindScriptObject数据文件前缀data/constants/opaque_obj_*OPAQUE_OBJ_FILENAME_PREFIX不透明自定义对象数据文件前缀data/sample_inputs/{}.ptSAMPLE_INPUTS_FILENAME_FORMAT样例输入data/executorch/EXECUTORCH_DIRExecuTorch 产物如 PTE 文件extra/EXTRA_DIR用户自定义附加文件data/mtiaMTIA_DIRMTIA 加速器产物需要说明的是archive_version的记录位于归档根目录常量值为archive_version而规范中.data/version一节所描述的是 zip 文件格式自身的版本源码注释明确区分了两者——PT2ArchiveWriter调用set_min_version(6)写入的是 zip 格式版本写向/.data/versionPT2 Archive 规范版本则写向/archive_version见 torch/export/pt2_archive/_package.py#L93-L117。三、Archive 头信息Archive Headersarchive_format声明归档使用的格式。目前只允许取值pt2。写入与校验都在源码中硬编码PT2ArchiveWriter.__exit__在关闭归档前若未写入该记录会自动补写ARCHIVE_FORMAT_VALUE而PT2ArchiveReader构造时会校验read_string(ARCHIVE_FORMAT_PATH) ! ARCHIVE_FORMAT_VALUE并抛出AssertionError(Invalid archive format)见 torch/export/pt2_archive/_package.py#L195-L211。byteorder取值为little或big供 zip 文件读取端判断字节序。/.data/version与/archive_version规范中/.data/version包含归档版本号并特别强调它既不是 export 序列化的 schema 版本也不是 ATen Opset 版本。从实现看PT2 Archive 规范自身的版本号记录在根目录archive_version文件中初始值为02024 年 9 月 4 日确立的初始版本。load_pt2加载时会读取archive_version并与当前ARCHIVE_VERSION_VALUE比对不一致即抛出ValueError见 torch/export/pt2_archive/_package.py#L1136-L1141。源码注释对版本号升级规则有明确约定当 PT2 Archive 规范出现破坏性变更例如删除某个文件夹、改变下列字段的命名约定时必须递增ARCHIVE_VERSION_VALUE且 Archive 读取端需要相应支持同时加载当前与旧版本的 PT2 Archive见 torch/csrc/export/pt2_archive_constants.h#L17-L24。/.data/serialization_id为当前归档生成的哈希用于校验。归档读取端PT2ArchiveReader同样暴露了archive_version()方法——当archive_version记录缺失时会将其视为版本 0从而保持对早期归档的向后兼容见 torch/export/pt2_archive/_package.py#L229-L240。四、AOTInductor 编译产物data/aotinductor/路径约定路径格式为/data/aotinductor/model_name-backend/。AOTInductor 编译产物按模型-后端组合分别保存例如model1模型在 A100 与 H100 上的编译产物会分别存放到model1-a100与model1-h100文件夹中互不干扰。文件夹典型内容该文件夹通常包含以下文件uuid.wrapper.so由uuid.cpp编译生成的动态库uuid.wrapper.cppAOTInductor 生成的 C wrapper 文件uuid.kernel.cppAOTInductor 生成的 C kernel 文件*.cubin由 Triton codegen kernel 编译出的 Triton kernel 二进制uuid.wrapper_metadata.json从aot_inductor.metadatainductor 配置传入的元数据可选uuid.json自定义算子需要由ProxyExecutor执行时的外部 fallback 节点按ExternKernelNode结构序列化。若模型未使用自定义算子/ProxyExecutor则该文件省略。打包实现细节在 torch/export/pt2_archive/_package.py#L254-L328 的_package_aoti_files中可以看到更精细的逻辑aoti_files可以传列表自动命名为model或{模型名: 文件列表}字典文件按 basename 判断去向以custom_obj_前缀开头的文件进入data/constants/其余文件进入data/aotinductor/model_name/源码注释特别指出早期实现拒绝一个模型存在多个.so的情况但Triton CPU AOTInductor 每个模型会产生多个.so文件wrapper、launcher、kernel因此当前实现支持多.so若存在权重会对权重做去重分组group_weights将共享同一 storage 的多个张量合并保存并在data/aotinductor/model_name/weights_config.json中记录weight_name - (filename, shape, stride, offset)的映射加载时通过as_strided还原见 torch/export/pt2_archive/_package.py#L1208-L1222。五、权重存储data/weights/存储约定路径为/data/weights/*。模型的参数parameters与缓冲区buffers保存在该文件夹中每个张量保存为独立文件文件内只有原始数据 blobraw data blob。张量的元数据以及模型权重 FQN → 原始数据 blob的映射单独保存在model_name_weights_config.json中。打包端两种序列化路径_package_state_dicttorch/export/pt2_archive/_package.py#L435-L479会遍历exported_program.state_dict把权重分成两类pickle 路径张量是 tensor subclass 且非 fake tensor 时_should_use_pickle见 L366-L367用torch.save保存整个张量对象文件名如weight_idxraw bytes 路径普通张量通过_get_raw_tensor_bytes直接拷贝untyped_storage的原始字节见 L339-L363并对共享 storage 的张量做去重合并。每条权重在 config 中记录为schema.PayloadMeta包含path_name、is_param是否是torch.nn.Parameter、use_pickle标志以及序列化的tensor_metashape/stride/storage_offset/dtype/device 等。加载端按 FQN 还原并保留别名关系_load_state_dicttorch/export/pt2_archive/_package.py#L864-L917先用_build_file_map把每个weight_*文件读成 flat 张量再对每个权重 FQN 用torch.as_strided(size, stride, storage_offset)从 flat 张量上切出正确的视图is_param为真的张量恢复为torch.nn.Parameter并还原requires_grad。源码注释指出这种FQN → 文件 → strided 张量的链式映射保留了权重之间的别名aliasing关系。另外加载端对旧格式做了兼容若归档中存在data/weights/model_name.pt这一历史遗留文件会走 legacy 分支读取并给出警告提示使用torch.export.save重新生成新格式见 L868-L875。六、常量存储data/constants/存储约定路径为/data/constants/*。三类对象保存在此Tensor 常量TensorConstants非持久缓冲区non-persistent buffersTorchBind 对象。元数据与模型常量 FQN → 原始数据 blob的映射保存在model_name_constants_config.json中。打包端三类前缀与两种格式_package_constantstorch/export/pt2_archive/_package.py#L482-L565按对象类型分类处理Tensor 常量torch.Tensor与权重类似subclass 走 pickle 路径普通张量走 raw bytes 路径文件名前缀为tensor_ScriptObjecttorch._C.ScriptObject使用 torchbind 感知的 picklerzip-archive 格式加载端可在 C 中读取文件名前缀为custom_obj_不透明自定义对象custom class object使用标准 Python pickle文件名前缀为opaque_obj_。源码注释特别说明不同的文件名前缀正是加载端区分格式的判别标志discriminator。加载端_load_constantstorch/export/pt2_archive/_package.py#L920-L985依据path_name的前缀分发tensor_前缀走与权重一致的 as_strided 还原或torch.loadcustom_obj_前缀用torch._C._pickle_load_obj反序列化opaque_obj_前缀用标准pickle.loads。若遇到未知前缀则抛出RuntimeError。同样地若归档中存在data/constants/model_name.pt旧格式文件也会走 legacy 兼容分支。七、样例输入data/sample_inputs/路径为/data/sample_inputs/model_name.pt。torch.export使用的样例输入sample input可以一并打包进归档供下游使用。通常它是一个扁平化的张量列表同时包含forward()的 args 与 kwargs实际打包时通过reorder_kwargs保证 kwargs 顺序与调用约定一致见 torch/export/pt2_archive/_package.py#L33。.pt文件由torch.save(sample_input)产生读取方式为Python 端torch.load()C 端torch::pickle_load()。当模型存在多份样例输入时会打包为model_name_index.pt如model1_0.pt、model1_1.pt。八、模型定义models/路径为/models/model_name.json。模型定义是torch.export.save产出的ExportedProgram 序列化 JSON外加其他模型级元数据。从打包流程看torch/export/pt2_archive/_package.py#L581-L626每个模型会依次写入data/weights/model_name_weights_config.json权重 payload 配置data/constants/model_name_constants_config.json常量 payload 配置models/model_name.jsonserialize()产出的 ExportedProgram 序列化数据此时serialize_state_dictFalse、serialize_constantsFalse即权重/常量不重复内嵌data/sample_inputs/model_name.pt样例输入。加载端_load_exported_programsL988-L1026扫描所有models/前缀条目从文件名反推模型名读取样例输入、反序列化 ExportedProgram再把上一步还原的state_dict与constants一并交给ExportedProgramDeserializer完成整体重建。九、多模型共存PT2 Archive 规范支持多个模型定义共存于同一个文件中。model_name是模型的唯一标识符并被归档其他文件夹作为引用依据——例如data/weights/model1_weights_config.json对应models/model1.jsondata/aotinductor/model1/对应model1的编译产物。在 Python API 层面传入dict[str, ExportedProgram]或dict[str, list[str]]AOTI 文件即可打包多个模型只传单个对象时自动命名为model见 torch/export/pt2_archive/_package.py#L590-L596。加载端返回的PT2ArchiveContents中exported_programs与aoti_runners均为以模型名为键的字典L788-L792。十、高层与底层 APIsave/load 与 package_pt2/load_pt2高层 APItorch.export.save与torch.export.loadtorch.export.savetorch/export/init.py#L266-L335内部就是对package_pt2的封装官方 docstring 给出了完整可运行示例import torch import io class MyModule(torch.nn.Module): def forward(self, x): return x 10 ep torch.export.export(MyModule(), (torch.randn(5),)) # 保存到文件 torch.export.save(ep, exported_program.pt2) # 保存到 io.BytesIO 缓冲区 buffer io.BytesIO() torch.export.save(ep, buffer) # 携带附加文件保存 extra_files {foo.txt: bbar.decode(utf-8)} torch.export.save(ep, exported_program.pt2, extra_filesextra_files)对应的torch.export.loadtorch/export/init.py#L338-L405import torch import io # 从文件加载 ExportedProgram ep torch.export.load(exported_program.pt2) # 从 io.BytesIO 对象加载 with open(exported_program.pt2, rb) as f: buffer io.BytesIO(f.read()) buffer.seek(0) ep torch.export.load(buffer) # 加载附加文件map 中的 value 会被替换为实际数据 extra_files {foo.txt: } ep torch.export.load(exported_program.pt2, extra_filesextra_files) print(extra_files[foo.txt]) print(ep(torch.randn(5)))需要注意两个官方 warning该 API 处于活跃开发阶段旧版本保存的文件可能无法在新版本中使用且torch.export.load底层使用 pickle切勿加载不可信来源的数据。底层 APIpackage_pt2与load_pt2package_pt2torch/export/pt2_archive/_package.py#L648-L728提供更细粒度的打包控制支持以下参数参数类型说明fstr / os.PathLike / IO[bytes]输出文件或可写缓冲区推荐以.pt2结尾exported_programsExportedProgram或dict[str, ExportedProgram]要保存的导出程序可多模型aoti_fileslist[str]或dict[str, list[str]]torch._inductor.aot_compile(..., {aot_inductor.package: True})生成的 AOTInductor 文件extra_filesdict[str, Any]附加文件写入extra/目录opset_versiondict[str, int]opset 名到版本号的映射pickle_protocolint覆盖默认 pickle 协议默认值为 2executorch_filesdict[str, bytes]可选的 ExecuTorch 产物若三者都未传入exported_programs、aoti_files、extra_files全为None会直接抛出断言错误。打包全程在PT2ArchiveWriter上下文管理器中完成关闭时自动补写archive_format与archive_version头记录。load_pt2torch/export/pt2_archive/_package.py#L1082-L1224的加载参数参数默认值说明f—文件路径或可读缓冲区expected_opset_versionNone期望的 opset 版本映射num_runners1加载 AOTInductor 产物的 runner 数量run_single_threadedFalse是否以单线程模式运行避免与 CUDA Graph 的同步逻辑冲突device_index-1加载目标设备索引CUDA 下-1对应cuda传1则加载到cuda:1load_weights_from_diskFalse是否把权重解出到磁盘再由 AOTI runner 加载加载后得到PT2ArchiveContentsexported_programs、aoti_runners、extra_files。其中aoti_runners里的每个元素是AOTICompiledModelL731-L785它是一个可调用对象调用时会取出AOTIModelPackageLoader的 call spec用pytree对输入输出做展平/还原然后执行boxed_run还提供了get_metadata()、get_constant_fqns()以及load_constants(constants_map, check_full_update..., user_managed..., allow_h2d_copy...)用于运行期更新常量。user_managedTrue时 loader 直接保存张量指针调用方必须保证其存活allow_h2d_copyTrue时 CPU 张量会被静默拷贝到模型所在设备便于把 CPU 的state_dict()加载进非 CPU 模型与user_managed不兼容。需要说明的是当传入的是内存缓冲区且归档内含 AOTI 产物时load_pt2会先把缓冲区内容写入临时.pt2文件再交给AOTIModelPackageLoader因后者不支持直接读缓冲区见 L1175-L1195。十一、版本与兼容性归档规范版本当前ARCHIVE_VERSION_VALUE为0后续任何破坏性变更删目录、改命名约定等都必须递增该值且读取端需要同时兼容新旧版本见 torch/csrc/export/pt2_archive_constants.h#L15-L24。load_pt2在版本不匹配时直接报ValueError。zip 格式版本PT2ArchiveWriter将 zip 最小版本设为 6该版本写入/.data/version与归档规范版本是两套独立体系。遗留格式兼容data/weights/model_name.pt与data/constants/model_name.pt属于历史遗留格式加载时会打印 warning 并提示改用torch.export.save重新生成新格式。读取端校验PT2ArchiveReader构造时校验archive_formatis_pt2_package可用来无损检测任意字节流/路径是否为 PT2 Archive适合在部署管线中做格式分发决策。总体而言PT2 Archive 通过头信息 按模型组织的分区目录 JSON 配置索引原始数据 blob的设计把ExportedProgram的图定义、样例输入、AOTInductor 编译产物、权重与常量统一收纳进一个 zip 文件既保持了标准 zip 的可操作性又通过model_name约定天然支持多模型、多后端产物共存是理解 PyTorch 2 系列导出与部署链路的关键规范。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
