什么是 torch.compile？Dynamo、AOT Autograd 与 Inductor 加速 PyTorch 详解 | IBM Technology 学习笔记

给 PyTorch 模型加一行 model = torch.compile(model)，就能变快吗？答案是：有机会，但不保证。这篇笔记整理自 IBM Technology 的视频，讲清楚 torch.compile 背后的三个阶段，以及“为什么没提速”时该问哪四个问题。

一行代码背后的“机会”

普通 PyTorch 执行时，Python 逐个分发操作；小操作的内核（kernel）启动开销和中间数据搬运会变得显著。torch.compile 让 PyTorch 有机会捕获一段计算、优化后生成更高效的代码。默认是惰性编译：第一次用输入调用模型时才开始追踪和编译，所以第一次运行可能更慢。

第一阶段：Dynamo 与守卫（guards）

Dynamo 追踪 Python 执行，把能表示的张量操作捕获成 FX 图，同时记录关于输入的假设，如张量形状、数据类型，这些就是守卫。之后每次调用先检查守卫，成立就复用编译结果，不成立可能重编译。

第二阶段：AOT Autograd

它为后端准备图：把高层操作分解成低层操作（一个操作可能变成好几个，所以不是“把图变小”）；训练时还会捕获反向计算，推理则没有反向图。

第三阶段：Inductor

Inductor 是默认后端编译器，负责融合兼容操作、安排调度、减少中间内存流量。比如“加偏置 + 激活函数”可融合成一个内核。GPU 上通常生成 Triton 内核，CPU 上生成 C++ 代码。

图中断、重编译与正确测速

• 图中断（graph break）：遇到无法捕获的代码，如依赖张量内部值的 Python 条件、.item()，图会被切开，不支持的部分用普通 Python 跑。程序仍正确，只是可优化区域变小。
• fullgraph=True 要求整体捕获成一张图，否则报错。
• 输入形状变化可能触发重编译；动态形状能减轻但不能消除。
• 测速要把编译时间和稳态执行时间分开，GPU 上还要做好同步。

没提速时问四个问题

1. 捕获的模型部分够多吗？
2. 图中断是否限制了优化？
3. 输入变化是否导致重编译？
4. 测的是编译成本还是稳态执行时间？

术语表

• Kernel（内核）：在 GPU 等硬件上运行的一段代码
• FX graph（FX 图）：Dynamo 捕获的操作图
• Guards（守卫）：编译结果可复用的前提假设
• Graph break（图中断）：无法捕获时把图切开
• Operator fusion（算子融合）：多个操作合并为一个内核
• Inductor：PyTorch 默认后端编译器

自测题

为什么加了 torch.compile 后第一次运行可能更慢？
A. 自动降低模型精度 B. 第一次调用时才开始追踪和编译 C. GPU 需要重启 D. Python 被替换成 C++
答案：B（编译是惰性的，首次调用要付出编译成本）

和求职面试的关系

做 AI/ML 工程岗，面试常问“如何加速 PyTorch 推理/训练”“什么是图中断”“怎么正确做基准测试”。能讲清 Dynamo、AOT Autograd、Inductor 的分工，并说出“编译时间与稳态时间分开测”，会明显加分。

来源

内容源自 IBM Technology 视频：What Is torch.compile? Dynamo, PyTorch & Faster AI Models Explained。完整学习卡与更多面试题：课程页。本文为学习笔记，非原视频转载。