你有一个能正常运行的 PyTorch 模型，并希望让它运行得更快。只需添加一行代码：model = torch.compile。第一次运行时可能会花费更长时间，因为 PyTorch 需要编译代码。此后，模型运行速度可能会提升。但 PyTorch 实际上在做什么？为什么编译对某些模型的加速效果远大于其他模型？要回答这个问题，我们需要了解三个阶段：Dynamo、AOT Autograd 和 Inductor。

在正常的 PyTorch 执行中，Python 会运行你的模型并在遇到操作时进行分发。某些操作已经调用高度优化的库，而其他操作可能需要单独启动内核（kernel）。内核是在 GPU 等硬件上运行的代码片段。对于小型操作而言，启动工作和移动中间数据产生的开销可能变得显著。Torch compile 让 PyTorch 有机会捕获一段计算过程，对其进行优化并生成执行效率更高的代码。这里的关键字是“机会”。编译并不能保证加速，也不一定能将整个模型转化为单个图结构。

你写的 model = torch.compile 这行代码，会在已有模型周围创建一个编译后的包装器。在标准设置下，编译是惰性的（lazy）。这意味着 PyTorch 会在你首次使用输入调用模型时开始追踪和编译。这次初始调用提供了张量形状、数据类型和设备等信息，PyTorch 利用这些信息对生成的代码进行特化。现在让我们看看在那次初始调用中发生了什么，第一阶段是 Dynamo。

Dynamo 追踪你的 Python 执行过程，并捕获它能表示为 FX 图的张量操作。FX 图是操作的表示形式。Dynamo 不仅追踪操作，还记录它们之间流动的信息。与简单存储原始 Python 源代码不同，Dynamo 还会记录关于输入和程序状态的假设。这些被称为守卫（guards）。例如，编译后的代码可能对特定的张量形状或数据类型有效。在后续调用中，PyTorch 会在重用编译代码之前检查这些假设是否仍然成立。

如果假设不再成立，PyTorch 可能会决定需要重新编译一个新版本。因此，Dynamo 的工作不仅仅是记录操作，它还确定哪些内容可以被捕获，以及在什么条件下可以重用这些被捕获的代码。第二阶段是 AOT Autograd（即时自动微分）。它的作用是为后端准备被捕获的计算过程。

这包括将一些高级操作分解为编译器可以处理的低级操作。最好将其视为准备图结构，而不仅仅是使其变小。例如，分解实际上可能将一个操作转化为多个操作。对于训练任务，AOT Autograd 还会捕获用于计算梯度的反向传播过程。它可以准备前向和反向图，以便编译器有机会对两者进行优化。对于不需要梯度的推理任务，则无需生成训练用的反向图。

此时，我们拥有了一个准备好供后端编译器使用的计算表示形式。第三阶段是 Inductor，它是 PyTorch 的默认后端编译器。Inductor 接收准备好的图结构并为目标硬件生成代码。它可以融合兼容的操作，选择工作调度方式并减少不必要的中间内存流量。例如，假设模型先添加偏置然后应用激活函数。编译器可能能够将这两个操作融合为一个内核启动，而不是分别启动内核并存储中间结果。

在 GPU 上，Inductor 通常生成 Triton 内核；在 CPU 上则生成 C++ 代码。当使用优化库（如矩阵乘法库）是更好的实现方式时，它也可以调用这些库。不过，并非所有操作都会被融合。大型矩阵乘法和其他专用操作可能仍会使用单独的内核库。目标是针对工作负载生成高效代码，而不是产生尽可能少的内核数量。

那么，当 Dynamo 遇到无法捕获的代码时会发生什么？在默认设置下，它可能会创建图断裂（graph break）。Dynamo 会固化目前已捕获的图并将其发送给后端，然后在常规 Python中运行不支持的代码并恢复追踪。例如，依赖于张量内部值的 Python 条件可能导致图断裂；调用 .item() 也可能导致这种情况。根据代码和编译器设置，程序可能仍能正确运行，但编译器现在需要优化的区域变小了。

这可能会减少融合机会并在边界处增加开销。但这并不意味着整个模型会失去加速效果。还有一个更严格的选项 fullgraph=True，它要求函数被捕获为单个图结构，并在无法做到时抛出错误。因此，当性能令人失望时，图断裂是首先值得检查的因素之一。编译本身也有成本。第一次调用承担了该成本，而后续调用如果守卫条件仍然匹配则可以重用相同的缓存编译代码。

如果输入形状或其他假设发生变化，PyTorch 可能需要重新编译。动态形状支持可以减少某些重新编译的情况，但并不能消除所有可能的特化需求。这就是为什么性能测量应将编译时间与稳态执行时间分开的原因。对于 GPU 测量，你还需要进行适当的同步，以确保测量的是实际完成的工作。

因此，在 torch.compile 背后要记住这三个阶段：Dynamo 捕获计算过程并记录其可重用的条件；AOT Autograd 准备图结构并为训练处理反向图的捕获；Inductor 为目标硬件生成代码。当编译没有帮助时，我会问四个问题：我们捕获了足够的模型部分吗？图断裂是否限制了优化？变化的输入是否导致了重新编译？我们测量的是编译成本还是稳态执行时间？

这些问题比仅仅知道模型已被编译能告诉你更多。torch.compile 的真正价值在于它让 PyTorch 有机会优化那些原本会单独执行的操作。这种机会是否能转化为加速，取决于模型和工作负载。