# torch.compile · 给动态图做编译

> eager 模式一个算子一个算子地即时执行，有 Python 开销、内核启动开销、还没法跨算子融合。torch.compile 怎么在不牺牲 Python 灵活性的前提下做图编译？看 TorchDynamo 抓字节码、装 guard、遇到搞不定的就「graph break」，以及 TorchInductor 把图编成 Triton/C++ 快内核。

- 作者：David（道雾轩）
- 专栏：PyTorch 源码漫游（https://daiw.org/manual/pytorch-internals.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/pytorch-internals/torch-compile
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

# torch.compile · 给动态图做编译

漫游的最后一站。前面看到的 PyTorch，是一台 **eager（即时执行）** 引擎：[一个算子经过 dispatcher 落到内核](https://daiw.org/manual/pytorch-internals/dispatch-journey)，算完，再来下一个。这很灵活、很好调试，但有性能天花板。`torch.compile`(PyTorch 2.0)就是来捅破这层天花板的——而它捅破的方式，相当巧妙。

## eager 模式的三笔开销

「一个算子一个算子地即时执行」，慢在哪？

1. **Python 开销**：每个算子都要绕一圈 Python 解释器 + [dispatcher](https://daiw.org/manual/pytorch-internals/dispatch-journey)，对小算子来说，这圈「手续」可能比计算本身还久；
2. **内核启动开销**：每个算子在 GPU 上是一次独立的 kernel 启动，启动本身有固定成本；
3. **没法跨算子优化**：`x.relu().add(1).mul(2)` 三个逐元素操作，eager 下是**三个内核、三趟显存往返**——明明可以**融合成一个内核、一趟往返**，但 eager 看不到「下一个算子」，融不了。

对大矩阵乘法，这些开销占比小；但对一堆小算子、或[受 Python 拖累喂不饱的 GPU](https://daiw.org/manual/pytorch-internals/cuda-and-memory)，它们能吃掉可观的性能。

## 难题：怎么给「动态 Python」做图编译

解决办法显而易见——**把一串算子捕获成一张图，整体优化（融合、去开销）、编译成快内核**。可 PyTorch 的命根子恰恰是**动态**：[define-by-run](https://daiw.org/manual/pytorch-internals/autograd-graph)，代码想怎么写就怎么写，有 `if`、有循环、有任意 Python。**怎么在不夺走这份灵活的前提下，把它编译了？**

这是老方案(`torch.jit.trace` / `script`)一直没做漂亮的地方：trace 抓不住控制流，script 要你把模型写成它能懂的子集，处处别扭。PyTorch 2.0 的答案是两个新部件：**TorchDynamo** 负责「抓图」，**TorchInductor** 负责「编译」。

## TorchDynamo：从字节码里抠出图，搞不定就「break」

**TorchDynamo** 是个 JIT，它的高招是**在 Python 字节码层面动手**：当一个被 `torch.compile` 包过的函数运行时，Dynamo 挂进 CPython 的帧求值，**分析这段函数的字节码**，把其中的 PyTorch 算子**抠出来，拼成一张 FX 图**，交给后端编译。

两个关键机制：

- **Guard（护栏）**：Dynamo 编出来的图，是**在某些前提下**才成立的——比如输入张量是某个 dtype、某个形状。它把这些前提记成 **guard**。下次调用，**guard 全过 → 直接复用编译好的图**（飞快）；guard 没过（比如换了输入形状）→ **重新编译**一份。
- **Graph break（图断裂）**：遇到 Dynamo **搞不定的东西**（任意 Python 副作用、依赖数据的控制流、不支持的操作），它不会崩，而是**优雅地「断一下」**：把搞不定的那段**退回 eager 执行**，然后**接着往下抓图**。

<Callout type="warn">
  **这解释了「torch.compile 有时飞快、有时没感觉」。** 如果你的 `forward` 干净、以张量运算为主，Dynamo 能抓出一大整张图，编译收益巨大；但如果里面塞了很多 `print`、`.item()`、依赖数值的 `if`、或古怪的 Python，就会**频繁 graph break**，图被切得稀碎，编译几乎没得优化。**调 `torch.compile` 性能，一大半工作就是「减少 graph break」**——可以用 `torch._dynamo` 的日志看它在哪儿断了。这种「搞不定就退回 eager」的设计，正是它能兼容任意 Python、又不像老方案那样动辄报错的关键。
</Callout>

## TorchInductor：把图编成 Triton / C++

抓到 FX 图后，默认后端 **TorchInductor** 负责把它**编译成真正快的内核**：

- **GPU 上，生成 [Triton](https://github.com/triton-lang/triton) 代码**（OpenAI 的 GPU 内核语言）；**CPU 上，生成 C++/OpenMP**。
- 核心优化是**融合(fusion)**：把上面那串 `relu → add → mul` 一类的逐元素操作**合并成一个内核**——一次启动、一趟显存往返，直接把 eager 的第 2、3 笔开销干掉。

按 PyTorch 2 论文的数据，在 A100 上、180 多个真实模型上，TorchInductor 拿到了 **2.27× 推理、1.41× 训练**的几何平均加速，且优于此前多个编译器。**「不牺牲 Python 灵活性的图编译」，这一次真的落地了。**

```mermaid
flowchart LR
  py["你的 forward<br/>(任意 Python)"] --> dyn["TorchDynamo<br/>抓字节码 → FX 图 + guard"]
  dyn -->|"能抓的部分"| ind["TorchInductor<br/>→ Triton(GPU)/ C++(CPU)"]
  dyn -.->|"graph break 的部分"| eager["退回 eager 执行"]
  ind --> fast["融合后的快内核"]
```

## 回到 nanoGPT：那一行，和那个 `_orig_mod.`

现在回看 [nanoGPT](https://daiw.org/manual/nanogpt/scaling-and-speed) 里那一行 `model = torch.compile(model)`，你彻底懂它了：它给模型套上 Dynamo，首次运行时抓图、编译（所以「要花约一分钟」），之后 guard 命中就跑编译好的快内核。

还有[那个当时留的疑问——`_orig_mod.` 前缀](https://daiw.org/manual/nanogpt/scaling-and-speed)：`torch.compile` 会用一个包装器裹住你的模型，原来的模型被存在 `_orig_mod` 这个属性下。于是编译后模型的参数名，前面都多了 `_orig_mod.`。[nanoGPT 断点续训时特意把这前缀剥掉](https://daiw.org/manual/nanogpt/train-setup)，就是为了让 checkpoint 的键名对得上没编译的模型。**一个上层包装的实现细节，在别处留下的那道「疤」，现在闭环了。**

<Callout type="info">
  一个漂亮的收尾观察：`torch.compile` **不取代**前面那台 eager 引擎，而是**架在它之上**。它先用 Dynamo 抓图，搞不定的部分**照样退回 eager**（[dispatcher + 内核 + autograd](https://daiw.org/manual/pytorch-internals/dispatch-journey) 那一整套）去跑。换句话说，**你这一路漫游看到的引擎，依然是地基；compile 只是在上面盖了一层「能优化就优化、不能就落回地基」的加速层。** 这也是它能既快又稳的原因。
</Callout>

## 漫游到头了

从 [Tensor 是一张说明书](https://daiw.org/manual/pytorch-internals/tensor-and-storage)，一路走到 `torch.compile` 把动态图编译成 Triton 内核——PyTorch 这台引擎的主要齿轮，你都亲眼见过了。最后一章，我们把整个技术栈拼成一张图，回望这趟从 Python 到金属的旅程。👉 [收尾：从 Python 到金属，整台引擎的全景](https://daiw.org/manual/pytorch-internals/recap)

## 参考文献

- [1] Ansel et al. “PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation.” ASPLOS 2024. ACM DOI: 10.1145/3620665.3640366. —— TorchDynamo、TorchInductor、torch.compile 的权威论文，本章主要依据。
- [2] Tillet et al. “Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations.” MAPL 2019. —— Inductor 的 GPU 后端语言。
- [3] PyTorch 文档：“torch.compile” 与 “TorchDynamo overview.” https://pytorch.org/docs/stable/torch.compiler.html
- 源码：`torch/_dynamo/`、`torch/_inductor/`(pytorch/pytorch)。