PyTorchtorch.compile().aot_compile()全解析提前编译、序列化产物与部署实战【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchaot_compile()是 PyTorch 在torch.compile之上提供的提前编译AOT, Ahead-of-Time接口与默认首次调用才编译的惰性流程不同它会把图追踪、Inductor 代码生成、Triton 内核编译与自动调优全部前移到编译阶段将结果打包成可落盘的序列化产物供生产环境冷启动、跨进程/跨机器部署以及基于 fake tensor 的交叉编译使用。本文以官方用户指南 docs/source/user_guide/torch_compiler/torch.compiler_aot_compile.md 为主线结合仓库源码与测试用例从 API 用法、序列化机制、后端选择、闭包/外部引用处理到分布式训练完整展开帮助你掌握一条编译一次、处处加载的 PyTorch 部署链路。⚠️实验性功能官方文档明确标注该特性处于实验阶段API 可能随版本调整请以当前仓库实现为准。一、什么是aot_compile()与标准torch.compile的差异标准torch.compile(fn)采用惰性编译函数在第一次真实调用时才被追踪、编译并缓存。而torch.compile(fn).aot_compile(example_inputs)则在编译期就完成图追踪graph tracing用示例输入example inputs对函数做一次完整的前向追踪Inductor 代码生成将追踪得到的 FX Graph 交给后端生成高性能内核Triton 内核编译与自动调优autotuning针对目标硬件编译/调优内核产物打包将编译结果、guard 状态、原始字节码、运行时环境等打包为可序列化的AOTCompiledFunction。从源码看这一整条链路集中在 torch/_dynamo/aot_compile.py 的aot_compile_fullgraph中实现。其内部通过convert_frame.fullgraph_capture(model, args, kwargs)完成一次性图捕获随后在torch._guards.tracing(...)与strict_autograd_cache、bundled_autograd_cache等 functorch 配置开启的上下文中调用 backend 生成可序列化可调用对象最后组装出CompileArtifacts见 torch/_dynamo/aot_compile.py#L45-L61。AOT 编译适合以下三类场景消除生产环境冷启动编译延迟把最耗时的编译/调优放在部署之前完成序列化编译产物跨进程、跨机器部署运行端不再重复编译交叉编译在宿主机上用 fake tensor 追踪为不同目标设备生成产物。二、aot_compile()与 AOTInductor 的定位区别维度AOTInductoraot_compile()输入对象torch.export导出的模型torch.compile包装的函数/模块产物形态共享库C 部署序列化 artifact加载回 Python 可调用对象运行环境非 Python 环境Python runtime典型场景C 服务端/移动端部署留在 Python 生态内、预计算编译两者参考实现分别在 torch/_inductor/aot_inductor.py 系列与 torch/_dynamo/aot_compile.py。选择原则很简单需要 C 部署选 AOTInductor希望留在 Python 中且要预计算编译选aot_compile()。三、快速上手3.1 编译一个自由函数free functionimport torch def fn(x, y): return x y # Step 1: AOT 编译。必须 fullgraphTrue不支持 graph break。 # example_inputs 是 (args_tuple, kwargs_dict) 二元组。 compiled_fn torch.compile(fn, fullgraphTrue).aot_compile( ((torch.randn(3, 4), torch.randn(3, 4)), {}) ) # Step 2: 运行编译后的函数。 result compiled_fn(torch.randn(3, 4), torch.randn(3, 4)) # Step 3: 保存产物到磁盘。 compiled_fn.save_compiled_function(compiled_add.pt) # Step 4: 在另一进程加载运行无需重新编译。 with open(compiled_add.pt, rb) as f: loaded_fn torch.compiler.load_compiled_function(f) result loaded_fn(torch.randn(3, 4), torch.randn(3, 4))几点实现层面的说明fullgraphTrue是硬性要求。在 torch/_dynamo/eval_frame.py#L1020-L1023 中若self.fullgraph为假会直接抛出Graph breaks are not supported with aot compile. Please use torch.compile(fullgraphTrue).缓存是必需的torch._inductor.config.force_disable_cachesTrue时aot_compile会直接报错见同文件 L1015-L1018保存采用原子写atomic_write_binary先写临时文件并fsync再os.replace落盘避免写入中断产生损坏文件torch/_dynamo/aot_compile.py#L177-L188。3.2 编译一个nn.Module编译模块时调用编译后模块的.forward.aot_compile(...)。由于forward的第一个参数是self编译产物在调用时需要把模块实例作为第一个参数传入import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self): super().__init__() self.linear nn.Linear(4, 4) def forward(self, x): return self.linear(x) model MyModel() # AOT 编译 forward 方法。 compiled_fn torch.compile( model, fullgraphTrue ).forward.aot_compile(((torch.randn(3, 4),), {})) # 以模块实例作为第一个参数调用。 result compiled_fn(model, torch.randn(3, 4)) # 因为模型参数 requires_grad反向传播可穿透编译函数。 loss result.sum() loss.backward() print(model.linear.weight.grad) # 梯度正确流回模型参数 # 保存与加载。 compiled_fn.save_compiled_function(compiled_model.pt) with open(compiled_model.pt, rb) as f: loaded_fn torch.compiler.load_compiled_function(f) # 从磁盘加载后反向依然可用。 model.zero_grad() result loaded_fn(model, torch.randn(3, 4)) result.sum().backward() print(model.linear.weight.grad)编译模块这一用法在测试 test/dynamo/test_aot_compile.py#L970-L973test_aot_compile_module中有对应覆盖。值得一提的还有super()调用场景test_aot_compile_with_super_calltest/dynamo/test_aot_compile.py#L1176-L1190验证了__class__作为自由变量被正确序列化、加载后调用结果与原始 eager 一致。四、API 参考4.1torch.compile(...).aot_compile(example_inputs)对torch.compile()包装的函数做提前编译。参数example_inputstuple[tuple[Any, ...], dict[str, Any]]——(args, kwargs)二元组作为追踪示例输入。它决定了产物生效的张量形状、dtype 与设备。返回值AOTCompiledFunction——行为与原函数一致、但执行预编译代码的可调用对象额外暴露save_compiled_function(path)——将编译产物序列化到磁盘disable_guard_check()——关闭运行时 guard 校验高级用法见下文guard 机制。前置要求必须给torch.compile()传fullgraphTrueAOT 编译不支持 graph breakbackend 必须可调用字符串 backend 如inductor、eager、aot_eager均受支持。4.2torch.compiler.load_compiled_function(file, *, f_globalsNone, external_dataNone)从文件加载之前保存的 AOT 编译函数。参数file——以二进制读模式打开的文件类对象内含序列化的编译函数f_globalsdict | None——编译函数的可选全局作用域。当原函数引用了用户自定义类型或其他非常规全局对象时必需external_datadict | None——加载到运行时环境的可选数据。当原函数捕获了无法序列化的对象如nn.Module实例时必需其键需与save_compiled_function(external_data...)传入的一致。返回值已从磁盘预加载编译结果的 callable。该入口实现在 torch/compiler/init.py#L988-L1015其核心是读取字节后调用AOTCompiledFunction.deserialize(data, f_globals, external_data)反序列化过程会以torch._inductor.config.patch(enable_autograd_for_aotTrue)包裹编译函数的还原保证训练语义一致torch/_dynamo/aot_compile.py#L303-L304。4.3 guard 机制与disable_guard_check()AOTCompiledFunction.__call__每次调用前都会执行 guard 校验形状/设备/dtype 与编译时示例输入不一致时会抛RuntimeErrortorch/_dynamo/aot_compile.py#L237-L244。disable_guard_check()可关闭该校验属于高级用法测试见test_aot_compile_disable_guard_checktest/dynamo/test_aot_compile.py#L803。另外编译时还会对 guard 做序列化安全过滤过滤全局变量与不支持的 guard 类型并保留 guard 状态与 guard manager跨进程加载后 guard 校验能力依旧完整。五、选择后端SerializableCallable接口aot_compile()可与任何实现了SerializableCallable接口的后端协同。该抽象定义于 torch/_dynamo/aot_compile_types.py#L71-L85要求实现serialize_compile_artifacts、deserialize_compile_artifacts与__call__三个成员。内置的inductor、eager、aot_eager后端开箱即用# 默认 inductor优化代码生成。 compiled_fn torch.compile(fn, fullgraphTrue, backendinductor).aot_compile( ((torch.randn(3, 4),), {}) ) # eager无代码生成便于调试。 compiled_fn torch.compile(fn, fullgraphTrue, backendeager).aot_compile( ((torch.randn(3, 4),), {}) )从源码看当使用 Inductor 或基于 AOTAutograd 的后端时编译结果会被包装成BundledAOTAutogradSerializableCallable再参与序列化torch/_dynamo/aot_compile.py#L399-L415若产物未实现SerializableCallable则直接报错提示该后端不兼容L417-L425。GraphModuleSerializableCallabletorch/_dynamo/aot_compile_types.py#L87-L134则是另一条纯 FX Graph 的序列化路径反序列化时在新建的FakeTensorMode中重建图模块并recompile()。序列化细节补充自定义 Triton 内核无法直接 pickle其 JITFunction 含不可序列化的_thread.RLock。aot_compile_types.py通过Triton Kernel Side Table机制在序列化时记录内核的(module_path, function_name)反序列化时按导入路径重新导入并恢复全局kernel_side_table见该文件头部注释 L43-L68。六、闭包与外部引用closures / external references6.1 闭包自由变量自动序列化捕获自由变量的函数闭包受支持闭包状态会随编译产物一起序列化scale 2 def fn(x, y): return (x y) * scale compiled_fn torch.compile(fn, fullgraphTrue).aot_compile( ((torch.randn(3, 4), torch.randn(3, 4)), {}) ) compiled_fn.save_compiled_function(scaled_add.pt) with open(scaled_add.pt, rb) as f: loaded_fn torch.compiler.load_compiled_function(f)闭包支持的底层实现AOTCompilePickler.reducer_override对 cell、code、module、绑定方法及嵌套函数分别注册了还原器torch/_dynamo/aot_compile.py#L119-L153运行时会从runtime_env中按co_freevars重建f_localsprepare_f_localsL197-L210。对应测试test_aot_compile_with_closure_save_and_loadtest/dynamo/test_aot_compile.py#L1157-L1174验证了闭包产物保存/加载后结果与原始函数一致。6.2f_globals为用户自定义类型提供命名空间当函数引用的用户自定义类型无法被反序列化器找到时用f_globals提供所需命名空间with open(my_fn.pt, rb) as f: loaded_fn torch.compiler.load_compiled_function( f, f_globalsmy_module.__dict__ )6.3external_data非可序列化对象的捕获当函数捕获了不可序列化对象如nn.Module实例时通过external_data显式注入# 保存。 compiled_fn.save_compiled_function( fn_with_model.pt, external_data{model: model}, ) # 加载。 with open(fn_with_model.pt, rb) as f: loaded_fn torch.compiler.load_compiled_function( f, external_data{model: model} )其原理AOTCompilePickler.persistent_id会把external_data中的对象映射为持久 ID若序列化过程中仍遇到其他nn.Module不在external_data中会收集进errors并在pickler.dump后抛出提示用户将这些对象标记为 external datatorch/_dynamo/aot_compile.py#L64-L80、L277-L281。反序列化端AOTCompileUnpickler.persistent_load若找不到对应键会给出明确的Missing required external reference to data错误L156-L169。捕获模块的完整用例见test_aot_compile_with_captured_moduletest/dynamo/test_aot_compile.py#L1342。七、训练支持aot_compile()开箱即用地支持训练只要有参数requires_grad编译会自动追踪joint forwardbackward 图将其分区并分别编译两个半图。得到的函数具备 autograd 感知——对其输出调用.backward()行为与预期一致。复用上文MyModelmodel MyModel() compiled_fn torch.compile( model, fullgraphTrue ).forward.aot_compile(((torch.randn(3, 4),), {})) # 使用 AOT 编译函数的训练循环。 optimizer torch.optim.SGD(model.parameters(), lr0.01) for _ in range(3): optimizer.zero_grad() output compiled_fn(model, torch.randn(3, 4)) loss output.sum() loss.backward() optimizer.step()保存/加载保持 autograd 支持——从磁盘加载后反向依然可用compiled_fn.save_compiled_function(train_model.pt) with open(train_model.pt, rb) as f: loaded_fn torch.compiler.load_compiled_function(f) output loaded_fn(model, torch.randn(3, 4)) output.sum().backward() # 梯度正确流动训练语义的序列化在实现上有专门处理序列化/反序列化均处于bundled_autograd_cacheTrue的 functorch 配置下torch/_dynamo/aot_compile_types.py#L177-L181、torch/_dynamo/aot_compile.py#L303-L304BundledAOTAutogradSerializableCallable本质上包装了 AOTAutograd 生成的serialize()结果。测试test_aot_module_simplified_serializable_autogradtest/dynamo/test_aot_compile.py#L973专门验证了序列化后的 autograd 行为。7.1 分布式训练DTensor compile_on_one_rank对使用DTensortorch.distributed.tensor.DTensor做张量并行的模型aot_compile()可与compile_on_one_rank组合产出与 rank 无关的编译图不带该标志mesh 坐标、shard 偏移等 rank 相关值会作为常量烘焙进编译图导致每个 rank 一张不同的图开启后这些值变为符号化在运行时计算所有 rank 共享同一产物。配置方式一torch.distributed.config.patchimport torch.distributed.config as dist_config with dist_config.patch(compile_on_one_rankTrue): compiled_fn torch.compile( model, fullgraphTrue ).forward.aot_compile(((example_input,), {}))配置方式二环境变量TORCH_DISTRIBUTED_COMPILE_ON_ONE_RANK1。该开关在 torch/distributed/config.py#L18-L21 中定义注意其 deprecation 提示新写法为torch.compiler.config.compile_on_one_rank。完整示例torchrun多卡张量并行训练# train_tp.py -- 运行方式: torchrun --nproc_per_node8 train_tp.py import torch import torch.distributed as dist import torch.distributed.config as dist_config import torch.nn as nn import torch.nn.functional as F from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor, Replicate from torch.distributed.tensor.parallel import ( ColwiseParallel, RowwiseParallel, parallelize_module, ) class FeedForward(nn.Module): def __init__(self, dim, hidden_dim): super().__init__() self.linear1 nn.Linear(dim, hidden_dim) self.linear2 nn.Linear(hidden_dim, dim) def forward(self, x): return self.linear2(F.relu(self.linear1(x))) def main(): dist.init_process_group(backendnccl) rank dist.get_rank() torch.cuda.set_device(rank) mesh init_device_mesh(cuda, (dist.get_world_size(),)) model FeedForward(64, 128).cuda() parallelize_module(model, mesh, { linear1: ColwiseParallel(), linear2: RowwiseParallel(), }) x DTensor.from_local( torch.randn(4, 64, devicefcuda:{rank}), mesh, [Replicate()], run_checkFalse, ) # 以 compile_on_one_rank 编译——所有 rank 得到相同图。 with dist_config.patch(compile_on_one_rankTrue): compiled_fn torch.compile( model, fullgraphTrue, ).forward.aot_compile(((x,), {})) # 训练循环。 optimizer torch.optim.Adam(model.parameters(), lr1e-3) for step in range(5): optimizer.zero_grad() x DTensor.from_local( torch.randn(4, 64, devicefcuda:{rank}), mesh, [Replicate()], run_checkFalse, ) out compiled_fn(model, x) loss out.to_local().sum() loss.backward() optimizer.step() if rank 0: print(fstep {step}: loss {loss.item():.4f}) dist.destroy_process_group() if __name__ __main__: main()保存/加载方式与单进程一致——在任意 rank 调用save_compiled_function在任意 rank 调用load_compiled_function。由于compile_on_one_rankTrue产出 rank 无关图同一份产物可在每个 rank 直接加载无需逐 rank 编译。八、局限性Limitations必须fullgraphTruetorch.compile一旦遇到 graph breakaot_compile()直接报错。测试test_aot_compile_graph_break_error_fmttest/dynamo/test_aot_compile.py#L885覆盖了错误信息格式输入形状被特化产物仅对 example inputs 给定的形状、dtype、设备有效不同形状的输入在运行时触发 guard 失败除非显式disable_guard_check()并非所有后端都支持自定义后端必须实现SerializableCallable接口才能兼容保存/加载。九、补充实践要点默认参数带默认参数的函数同样支持编译与序列化见test_aot_compile_with_default_argstest/dynamo/test_aot_compile.py#L1204-L1216全局张量引用函数内引用模块级张量如EPS在 eager 与编译产物间行为一致test_aot_compile_with_global_tensortest/dynamo/test_aot_compile.py#L1192-L1202source_info 溯源AOTCompiledFunction.source_info()返回追踪到的源码信息SourceInfo便于调试定位编译来源交叉编译借助 fake tensor 的跨设备追踪test_cross_aot_compile、test_cross_compile_realistic_transformer_modeltest/dynamo/test_aot_compile.py#L1494可实现宿主机为其他目标设备编译。参考链接官方用户指南docs/source/user_guide/torch_compiler/torch.compiler_aot_compile.md核心实现torch/_dynamo/aot_compile.py、torch/_dynamo/aot_compile_types.py入口封装torch/_dynamo/eval_frame.pyaot_compile方法、torch/compiler/init.pyload_compiled_function分布式开关torch/distributed/config.py测试覆盖test/dynamo/test_aot_compile.py、test/distributed/tensor/test_dtensor_compile.py【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考