PyTorch Compile 2：前端

Python 太动态，图怎么抓#

要「抓图」，得先明确图是什么。这里的图是计算图：节点是算子、边是张量数据流，脱离了 Python 的逐行控制流，才能整体交给后端做融合、调度、代码生成。问题是 Python 代码里并不存在这样一张现成的图，它只是一段可能带分支、循环、副作用的普通程序。

举个具体例子。下面这个函数做了三步张量运算：

1
def f(x, w, b):
2
    y = x @ w             # 矩阵乘
3
    y = y + b             # 加偏置
4
    return torch.relu(y)  # 激活

把它对应的计算图画出来是这样：x、w、b 是输入，matmul、add、relu 三个算子按依赖顺序相连，每个算子算出的中间张量沿边流入下一个算子，最后一个算子的输出就是返回值。

这张图只记录「哪些张量运算、按什么依赖顺序发生」，而不关心它们写成一行还是三行、中间变量叫 y 还是别的名字——图里没有 Python 变量名，也没有语句的先后行号，只有算子和张量之间的数据依赖。「抓图」要做的，就是从这段 Python 代码里把这层结构识别出来、记下来，大致得到一串按依赖排好序的算子：

1
输入：x, w, b
2
  t1 = matmul(x, w)
3
  t2 = add(t1, b)
4
  t3 = relu(t2)
5
输出：t3

反过来，与张量无关的东西不会成为图上的节点。比如函数里夹一句 print(...)，或用一个普通 Python 整数记录调用次数，它们既不是张量运算、也不改变张量之间的数据流，自然不出现在计算图里。正因为图里只剩张量运算和它们的依赖关系，后端才能把它当成一个整体来做融合、调度和代码生成——这是 eager 模式下一行行单独下发算子时做不到的。

前一代方案 TorchScript 走的是「提前把整段 Python 静态分析成图」的路子：torch.jit.script 解析 Python 源码、但只支持一个受限子集，稍微动态一点的写法就不认识，用户得改代码去迁就它；torch.jit.trace 则跑一遍记录算子，只留下跑过的那一条路径，数据依赖的控制流会被拍平成某一次的结果，换个输入就错。两者共同的问题是：都想一次性把整个程序变成图，而真实模型里总有一些无法静态化的部分，遇到就卡住。

Dynamo 换了个思路：不强求整段都能变成图，而是能抓多少抓多少——把能识别的 PyTorch 运算抽成图交给后端，抓不动的地方就退回 Python 自己执行。这种「部分捕获」（partial graph capture）的能力，是它能在几乎不改代码的前提下适配真实模型的关键。

挂进 CPython：帧求值钩子（PEP 523）#

Dynamo 是怎么在函数运行时「拦下」它的？靠的是 CPython 在 PEP 523 里开的一个口子：帧求值 API（frame evaluation API）。

CPython 执行一次 Python 函数调用时，会为这次调用创建一个帧对象（frame，装着字节码、局部变量、运行栈等），再交给一个默认的求值函数逐条解释里面的字节码。PEP 523 允许把这个求值函数替换掉——它在解释器状态上放了一个函数指针，第三方可以改写它，从而接管「如何执行一个帧」。

Dynamo 就是把这个指针换成自己的钩子。之后每当一个被 torch.compile 标记过的函数要执行，控制权先落到 Dynamo 手里：它拿到帧的字节码，先不真正运行，而是尝试分析它、把里面的 PyTorch 运算抽成图并编译好，再改写帧的字节码让它去调用编译产物，最后把改写后的帧交回解释器执行。对用户来说，函数的调用方式没有任何变化，介入完全发生在解释器内部。

选在字节码这一层介入是有讲究的：字节码比 Python 源码规整得多（指令种类有限、栈式语义明确），不受源码写法花样的干扰，又比拿到具体数值的运行结果早——正好卡在「已经知道要做什么、但还没真做」的时机上，适合抓图。

符号执行字节码，抓出 FX graph#

拿到帧之后，Dynamo 做的不是真的执行字节码，而是符号执行（symbolic execution）：它模拟一个 Python 解释器，逐条「走」字节码，但栈上和变量里放的不是真实值，而是对这些值的符号化追踪对象。张量被表示成图里的一个节点，对张量做的每个 PyTorch 运算都被记成图上的一个算子；而与张量无关的普通 Python 运算（整数加减、构造列表、访问属性等）则在符号执行时就地算掉或如实建模，不进图。

抓出来的图是一张 FX graph（torch.fx 的图表示），节点是算子、边是张量流动，是一段脱离了 Python 控制流、可以整体优化的中间表示。这张图连同它的输入输出接口，被交给后端（默认 Inductor，下一篇的主题）去编译。

关键在于符号执行遇到「看不懂」的东西时会怎样。如果字节码里出现 Dynamo 无法符号化的操作——典型的是依赖张量具体数值的控制流，或调用一个它不认识、无法追踪的外部函数——它没法把这一段变成图上的节点。这时它不会崩溃，而是触发 graph break。

Graph break：抓不动就切开#

Graph break（计算图断裂）是 Dynamo「部分捕获」策略的直接体现。当符号执行走到一个抓不动的点，Dynamo 把这段函数从这里切开：

• 断点之前已经追踪到的算子，收尾成一张图，编译；
• 断点处那段抓不动的代码，原样交给 CPython 用 eager 模式执行；
• 断点之后的代码，另起一段继续符号执行、抓成第二张图，编译。

实现上，Dynamo 会为断点后的部分生成一个「续跑函数」（continuation function，形如 resume_at_X），把 eager 执行的结果接进去，让整个流程无缝衔接。结果是：程序行为和原来完全一致，但一次调用可能对应多张编译图，中间夹着若干段 eager。

Graph break 不影响正确性，却直接影响性能，原因有二：一是被切碎的图各自太小，后端能做的跨算子融合、调度优化都被局限在段内，断点两侧的算子无法融到一起；二是每次断裂都要付出图与 eager 之间来回切换的开销。断点一多，加速就被吃掉大半，极端情况下甚至比不编译还慢。所以调优 torch.compile 的很大一部分工作，就是找出并消除不必要的 graph break——第一篇提到的 fullgraph=True，正是用来把它们全部暴露成报错的开关。

常见的触发源包括：依赖张量值的 if/while（值要到运行时才知道，符号执行无法决定走哪支）、.item()/.tolist() 这类把张量转成 Python 数值的操作、print 和其他有外部副作用的调用，以及 Dynamo 尚未支持的第三方库函数。

Guard：复用编译产物的前提#

抓图和编译都发生在运行时，代价不小。要让这份代价摊薄到之后每次调用，就得复用编译产物。但复用不能是无条件的：这次抓出来的图，其实隐含了一堆「当时成立的前提」——输入张量是这个 dtype、这个 shape、在这个 device 上，某个 Python 标志位是这个值，某个对象是那个对象……换一组前提，同一份编译代码就未必对。

Dynamo 把这些前提显式记录下来，就是 guard（守卫）。抓图时它一边追踪一边收集依赖，编译产物上因此挂着一组 guard；论文里提到当时约有 30 种 guard，涵盖张量的各种属性（dtype、shape 等）、Python 变量的类型、常量特化、对象属性、dict/list/tuple 的结构、nn.Module 实例，以及全局的 PyTorch 状态等。

有了 guard，复用就变成一次快速检查：

• 函数被调用 → 先跑 guard 检查；
• 全部通过（命中）→ 直接运行已编译好的代码，跳过抓图和编译；
• 有 guard 不通过（未命中）→ 这份产物不能用，触发重新编译。

guard 检查在每次调用时都要先跑一遍，必须足够快，否则这份固定开销会侵蚀掉复用编译产物换来的加速，所以它被实现得很轻量。

Recompile：guard 失效就重编#

当一次调用没能通过已有的 guard，Dynamo 就为这组新前提重新抓图、重新编译，产出另一份带着新 guard 的产物，和旧的一起缓存起来。之后每次调用依次比对这些缓存条目，命中哪份用哪份。偶尔重编一两次没关系，只是一次性成本；真正的麻烦是反复重编。

最常见的重编触发源就是输入 shape 在变：如果每个 batch 的 shape 都不一样，而又没开动态 shape（下一节），那么每种新 shape 都会 miss guard、触发一次重编，缓存条目越堆越多。为防止这种「编译爆炸」拖垮程序，Dynamo 设了上限：torch._dynamo.config.cache_size_limit（新名字 recompile_limit，默认 8）限制同一个函数最多缓存多少份编译条目，accumulated_cache_size_limit（accumulated_recompile_limit，默认 256）是全局累计上限。撞到上限后，Dynamo 会对这个函数放弃编译、直接退回 eager，并打日志提醒。

所以看到「明明编译了却没加速」，除了 graph break，另一个高频原因就是悄悄发生的反复重编——它往往不报错，只是默默把编译产物一次次作废。

动态 shape：一份图吃多种尺寸#

对付「shape 老在变」的正道不是加大缓存上限，而是让一份图能同时服务多种 shape，这就是动态 shape（dynamic shapes）。

默认情况下 Dynamo 是「按具体 shape 特化」的：抓图时把 shape 当成常量写死进图（比如认定第 0 维就是 32），这样后端能利用确定的尺寸做更激进的优化，但 shape 一变就得重编。动态 shape 的做法是把某些维度当成一个符号变量——PyTorch 里叫 SymInt（符号整数）——图里对这个维度的运算都以符号形式记录，于是同一份图 [s0, 3] 可以同时吃 (2, 3)、(4, 3)、(8, 3) 等多种输入，不必为每个尺寸各编一份。

dynamic 参数控制这套机制的三种态度：

• dynamic=False：强制静态特化，每种 shape 各编一份，最快但最容易重编；
• dynamic=True：一开始就尽量把维度做成符号的，用一份动态图覆盖多种 shape；
• dynamic=None（默认）：自动——先按具体 shape 编（当静态处理），一旦发现某个维度在不同调用间变化、触发了重编，就把那个维度改成动态的重编一次，之后不再为它反复重编。

自动模式是个务实的折中：既让稳定不变的 shape 享受静态特化的性能，又不让偶尔变化的维度陷入无休止的重编。

定位问题的调试工具#

前面两类性能杀手——graph break 和反复 recompile——都不报错，得靠工具主动观察。常用的有这么几样：

• TORCH_LOGS 环境变量：最顺手的入口。TORCH_LOGS="graph_breaks" 打印每一次 graph break 及其原因和位置，TORCH_LOGS="recompiles" 打印每次重编及触发它的 guard，recompiles_verbose 会列出所有失败的 guard 检查，guards 则打印每份编译产物挂的 guard；可以逗号连写，如 TORCH_LOGS="graph_breaks,recompiles"。
• torch._dynamo.explain()：对一个函数跑一遍 Dynamo，汇总它一共抓出几张图、发生了哪些 graph break 及原因，适合快速体检某段代码「碎不碎」。
• tlparse：配合 TORCH_TRACE 记录的编译日志，把整个编译过程解析成结构化的 HTML 报告，适合排查大模型里成规模的编译行为。
• depyf：把 Dynamo 改写出来的字节码反编译回可读的 Python，让你直接看到「Dynamo 到底把我的代码变成了什么样」，也支持配合断点调试。

排查顺序上，一般先用 TORCH_LOGS="graph_breaks,recompiles" 跑一遍看有没有异常，定位到可疑处再用 explain 或 depyf 深入。

小结#

Dynamo 是 torch.compile 的前端，负责「抓图」这一步，整条链路可以串成一句：用 PEP 523 的帧求值钩子挂进 CPython，在函数执行前拦下它的字节码；对字节码做符号执行，把 PyTorch 运算抽成一张 FX graph 交给后端，抓不动的地方触发 graph break、退回 eager；抓到的图挂上 guard 记录复用前提，下次调用查 guard，命中就复用、未命中就 recompile；shape 变化则用 SymInt 撑起的动态 shape 来避免反复重编。理解了这套机制，第一篇提到的两个坑就都有了着落：graph break 来自「抓不动」，recompile 来自「guard 失效」。

图抓出来之后，真正把它变成快代码是后端的活——AOTAutograd 补上反向图、算子降解、Inductor 做融合与代码生成——那是下一篇的内容。