PyTorch Compile 1：基础用法

一行代码的承诺#

对一个已有模型，加速只需要一次包裹：

1
import torch
2

3
model = MyModel().cuda()
4
compiled = torch.compile(model)
5

6
output = compiled(x)  # 调用方式和原模型完全一样

compiled 的调用方式、输入输出、求导行为都和 model 一致，可以直接替换进已有的训练或推理循环。加速主要来自两处：减少 Python 逐算子下发的开销，以及把多个算子融合、减少对显存的反复读写。加速幅度取决于模型、硬件和 batch 大小，对计算规整、Python 开销占比高的模型收益更明显。

它在背后做了什么#

理解 torch.compile 最重要的一点：它不是源码转译，也不会把 Python 改写成另一种语言。它的工作方式是运行时抓图再编译，可以拆成三步。

1. 抓图（capture）：第一次调用被编译的函数时，torch.compile 追踪这段 Python 的执行，识别其中的 PyTorch 算子，把它们记录成一张计算图（computation graph）。tensor 运算会进图，与 tensor 无关的普通 Python 逻辑不会。
2. 编译（compile）：这张图交给后端编译器（默认是 TorchInductor），生成针对具体设备优化过的代码——GPU 上是 Triton kernel，CPU 上是 C++/OpenMP。
3. 缓存复用（guard + cache）：编译产物不是无条件复用的。抓图时 torch.compile 记下这次编译依赖的前提（输入的 dtype、shape、某些标志位等），这些前提叫 guard。之后每次调用先检查 guard，命中就直接运行已编译好的代码，不命中就重新编译。

整条编译栈从左到右是：TorchDynamo 负责前端抓图，AOTAutograd 把反向传播也一并抓进来（训练需要反向图），PrimTorch 把种类繁多的算子降解成一小撮基础算子，TorchInductor 作为后端把图编译成最终的 kernel。这张图后面几篇会分别深入。

由此可以直接得到两条使用时最需要记住的结论。

• 它不改变语义。编译后的结果和 eager 模式一致（数值上可能有正常范围内的浮点误差）。抓图时遇到实在处理不了的 Python 代码，torch.compile 不会硬编译或崩溃，而是退回 eager 执行那一段、再接着编译后面——这就是下面要讲的 graph break。
• 它有预热成本。抓图和编译都发生在运行时，集中在第一次调用（以及每次触发重新编译时）。所以第一次调用往往明显更慢，之后才快；计时时要跳过预热的前几次调用。

基本用法#

torch.compile 既可以作用在模型（nn.Module）上，也可以作用在普通函数上。最常见的是包裹后返回编译版本：

1
compiled_model = torch.compile(model)
2

3
@torch.compile
4
def forward(x, w):
5
    return torch.relu(x @ w)

nn.Module 还有一个原地编译的写法，编译自身而不返回新对象：

1
model.compile()

几种写法效果一致。被编译的对象照常参与训练：反向传播、优化器、混合精度都能正常工作，因为 AOTAutograd 会把反向图一起抓出来。唯一要留意的仍是预热——第一次调用包含抓图和编译的耗时，可能几秒到几十秒（max-autotune 更久），但这是一次性的。

三个 mode 的取舍#

torch.compile 用 mode 参数在「编译花多少功夫」和「跑得多快」之间选点：

1
compiled = torch.compile(model, mode="max-autotune")

主要有三档：

• default：默认档，在运行性能和编译开销之间取平衡，编译快、加速稳妥，适合大多数情况和日常开发。
• reduce-overhead：在此基础上叠加 CUDA graph，进一步压掉 Python 下发和 kernel 启动的开销。对小 batch、小模型这种「单个 kernel 很快、启动开销占比高」的场景收益明显，代价是多占一些显存。
• max-autotune：最激进的一档。编译时为关键算子（如矩阵乘、卷积）实测多种 kernel 实现并选最快的，同时默认开启 CUDA graph。编译明显更慢，但运行最快，适合要长期反复运行、值得为编译多花时间的模型。此外还有一个 max-autotune-no-cudagraphs 变体，做同样的 autotuning 但不启用 CUDA graph，用于 CUDA graph 会带来问题的情形。

经验上，开发和大多数场景先用 default；上线前想再压性能、且能接受更长编译时间，再按场景试 reduce-overhead（尤其小 batch）或 max-autotune。

几个关键旋钮#

除了 mode，还有三个参数会频繁遇到，机制留到后面几篇展开。

• fullgraph（默认 False）：控制对 graph break 的态度。默认允许 graph break，抓不动就切开、退回 eager；设为 True 则要求整段代码抓成一张完整的图，一旦遇到 graph break 就直接报错。开发时常用 fullgraph=True 主动暴露并逐个消除 graph break。
• dynamic（默认 None）：控制对输入 shape 的处理。False 为每种具体 shape 各特化一份（最快，但 shape 一变就重编）；True 从一开始就尽量编译出能吃多种 shape 的版本，避免反复重编；None 是自动——先按具体 shape 编，发现 shape 在变时再转成动态。
• backend（默认 "inductor"）：指定后端编译器。TorchInductor 是默认且通用的后端，也可以换成用于调试的 eager、aot_eager，或接入自定义后端。

两个必须知道的坑#

用 torch.compile 时，性能不及预期几乎总能归到这两件事上。

Graph break（计算图断裂）：抓图时遇到无法追踪的 Python 代码（例如依赖具体数值的控制流、print、调用某些不可追踪的库），torch.compile 会把图从这里切开——断点前后各编译成一张图，中间那段用 eager 运行，再把结果接起来。程序仍然正确，但每断一次就少一块可优化区域，断点两侧也无法融合；graph break 太多，加速就会被吃掉大半。

Recompile（重新编译）：编译产物靠 guard 复用。当某次调用不满足已有 guard——最常见的是输入 shape 变了、dtype 变了，或某个 Python 标志位变了——torch.compile 会为新情况另编一份。偶尔重编无所谓，但如果每次调用 shape 都在变而又没开动态 shape，就会反复重编，甚至触及重编次数上限，让编译开销盖过收益。

这两件事的定位手段和内部机制，是下一篇的主题。

小结#

torch.compile 的心智模型可以压成一句话：运行时把 PyTorch 代码抓成计算图、编译成优化过的 kernel，再靠 guard 判断能否复用；抓不动的地方退回 eager（graph break），前提变了就重编（recompile）。用法上，torch.compile(model) 一行包裹即可，mode 决定编译投入多少、换取多少运行性能，fullgraph、dynamic、backend 是最常调的旋钮。