torch.compile · 给动态图做编译
eager 模式一个算子一个算子地即时执行，有 Python 开销、内核启动开销、还没法跨算子融合。torch.compile 怎么在不牺牲 Python 灵活性的前提下做图编译？看 TorchDynamo 抓字节码、装 guard、遇到搞不定的就「graph break」，以及 TorchInductor 把图编成 Triton/C++ 快内核。
torch.compile · 给动态图做编译
漫游的最后一站。前面看到的 PyTorch，是一台 eager（即时执行） 引擎：一个算子经过 dispatcher 落到内核，算完，再来下一个。这很灵活、很好调试，但有性能天花板。torch.compile(PyTorch 2.0)就是来捅破这层天花板的——而它捅破的方式，相当巧妙。
eager 模式的三笔开销
「一个算子一个算子地即时执行」，慢在哪？
- Python 开销：每个算子都要绕一圈 Python 解释器 + dispatcher，对小算子来说，这圈「手续」可能比计算本身还久；
- 内核启动开销：每个算子在 GPU 上是一次独立的 kernel 启动，启动本身有固定成本；
- 没法跨算子优化：
x.relu().add(1).mul(2)三个逐元素操作，eager 下是三个内核、三趟显存往返——明明可以融合成一个内核、一趟往返，但 eager 看不到「下一个算子」，融不了。
对大矩阵乘法，这些开销占比小；但对一堆小算子、或受 Python 拖累喂不饱的 GPU，它们能吃掉可观的性能。
难题：怎么给「动态 Python」做图编译
解决办法显而易见——把一串算子捕获成一张图，整体优化（融合、去开销）、编译成快内核。可 PyTorch 的命根子恰恰是动态：define-by-run，代码想怎么写就怎么写，有 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 执行，然后接着往下抓图。
这解释了「torch.compile 有时飞快、有时没感觉」。 如果你的 forward 干净、以张量运算为主，Dynamo 能抓出一大整张图，编译收益巨大；但如果里面塞了很多 print、.item()、依赖数值的 if、或古怪的 Python，就会频繁 graph break，图被切得稀碎，编译几乎没得优化。调 torch.compile 性能，一大半工作就是「减少 graph break」——可以用 torch._dynamo 的日志看它在哪儿断了。这种「搞不定就退回 eager」的设计，正是它能兼容任意 Python、又不像老方案那样动辄报错的关键。
TorchInductor：把图编成 Triton / C++
抓到 FX 图后，默认后端 TorchInductor 负责把它编译成真正快的内核：
- GPU 上，生成 Triton 代码（OpenAI 的 GPU 内核语言）；CPU 上，生成 C++/OpenMP。
- 核心优化是融合(fusion)：把上面那串
relu → add → mul一类的逐元素操作合并成一个内核——一次启动、一趟显存往返，直接把 eager 的第 2、3 笔开销干掉。
按 PyTorch 2 论文的数据，在 A100 上、180 多个真实模型上，TorchInductor 拿到了 2.27× 推理、1.41× 训练的几何平均加速，且优于此前多个编译器。「不牺牲 Python 灵活性的图编译」，这一次真的落地了。
回到 nanoGPT：那一行，和那个 _orig_mod.
现在回看 nanoGPT 里那一行 model = torch.compile(model)，你彻底懂它了：它给模型套上 Dynamo，首次运行时抓图、编译（所以「要花约一分钟」），之后 guard 命中就跑编译好的快内核。
还有那个当时留的疑问——_orig_mod. 前缀：torch.compile 会用一个包装器裹住你的模型，原来的模型被存在 _orig_mod 这个属性下。于是编译后模型的参数名，前面都多了 _orig_mod.。nanoGPT 断点续训时特意把这前缀剥掉，就是为了让 checkpoint 的键名对得上没编译的模型。一个上层包装的实现细节，在别处留下的那道「疤」，现在闭环了。
一个漂亮的收尾观察：torch.compile 不取代前面那台 eager 引擎，而是架在它之上。它先用 Dynamo 抓图，搞不定的部分照样退回 eager（dispatcher + 内核 + autograd 那一整套）去跑。换句话说，你这一路漫游看到的引擎，依然是地基；compile 只是在上面盖了一层「能优化就优化、不能就落回地基」的加速层。 这也是它能既快又稳的原因。
漫游到头了
从 Tensor 是一张说明书，一路走到 torch.compile 把动态图编译成 Triton 内核——PyTorch 这台引擎的主要齿轮，你都亲眼见过了。最后一章，我们把整个技术栈拼成一张图，回望这趟从 Python 到金属的旅程。👉 收尾：从 Python 到金属，整台引擎的全景
参考文献
- [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)。