Pytorch 学习笔记(9): PyTorch.Compile

一、概述torch.compiler是 PyTorch 2.0 引入的编译器模块提供了模型编译、优化和部署的核心 API。本文系统整理所有官方 API 及其用途。二、核心编译 API1.torch.compiler.compile功能编译 PyTorch 模型/函数说明详见torch.compile()的参数文档这是最常用的编译入口。# 典型用法modeltorch.compile(model,backendinductor)2.torch.compiler.reset功能重置编译器状态说明清除所有编译缓存将系统恢复到初始状态。适用于调试内存泄漏或状态异常问题。3.torch.compiler.disable功能禁用编译的装饰器说明为特定函数添加装饰器强制跳过编译。torch.compiler.disabledefmy_function(x):# 此函数不会被编译returnx*24.torch.compiler.set_stance功能设置编译器立场说明控制编译器的行为模式如 eager、compile 等模式切换。三、图操作与内联控制5.torch.compiler.allow_in_graph功能允许函数直接进入计算图说明告诉 Dynamo 前端跳过符号内省直接将函数写入图中。适用于 C 扩展函数等无法被追踪的代码。6.torch.compiler.substitute_in_graph功能注册 polyfill 处理器说明为 C 扩展函数提供替代实现在图内联时使用自定义版本替换原函数。7.torch.compiler.nested_compile_region功能标记嵌套编译区域说明标识一组可重复使用的操作编译一次后安全复用。适用于模型中重复出现的子结构。withtorch.compiler.nested_compile_region():# 这部分代码会被识别为可复用区域xself.block(x)四、常量与假设优化8.torch.compiler.assume_constant_result功能标记函数返回常量说明告知编译器某函数的输出是常量可进行常量折叠优化。五、后端与配置9.torch.compiler.list_backends功能列出可用后端说明返回所有有效的后端名称字符串如inductor、aot_eager等。print(torch.compiler.list_backends())# [inductor, cudagraphs, aot_eager, ...]10.torch.compiler.config功能编译器配置说明访问和修改编译器的全局配置参数。六、状态检测 APIAPI功能说明is_compiling()检测当前是否处于torch.compile()或torch.export()的编译/追踪过程中is_dynamo_compiling()检测是否通过 TorchDynamo 进行图追踪is_exporting()检测是否处于模型导出流程中典型用法iftorch.compiler.is_compiling():# 编译时优化路径passelse:# eager 模式路径pass七、分布式与 CUDA Graphs11.torch.compiler.set_enable_guard_collectives功能启用 guard 评估中的集合通信说明允许在 guard 评估时使用 collectives 同步多卡行为。12.torch.compiler.cudagraph_mark_step_begin功能标记 CUDA Graph 步骤开始说明指示新的推理/训练迭代即将开始用于 CUDA Graph 捕获优化。八、Guard 控制⚠️ 不安全 API以下 API 均以_unsafe结尾用于高级调优可能影响正确性API功能keep_portable_guards_unsafe()仅保留 Python/非 Python 环境通用的 guardkeep_tensor_guards_unsafe()保留所有张量 guardskip_guard_on_inbuilt_nn_modules_unsafe()跳过内置 nn 模块如nn.Linear的 guardskip_guard_on_all_nn_modules_unsafe()跳过所有 nn 模块的 guardskip_guard_on_globals_unsafe()跳过全局变量的 guardskip_all_guards_unsafe()跳过所有 guard⚠️警告这些 API 会放宽编译器的正确性检查仅在确信代码行为不变时使用。九、快速参考表场景推荐 API编译模型torch.compile()调试编译问题torch.compiler.reset()排除特定函数torch.compiler.disableC 扩展函数兼容torch.compiler.allow_in_graph()检测编译状态torch.compiler.is_compiling()查看可用后端torch.compiler.list_backends()CUDA Graph 优化torch.compiler.cudagraph_mark_step_begin()十、相关资源Torch Compile 官方文档TorchDynamo 细粒度追踪 APITorchInductor 溯源追踪提示本文基于 PyTorch 2.11 稳定版文档整理建议结合官方最新文档使用。