★ FEATURED ARTICLE 接手过几个大模型的推理优化项目之后我对 PyTorch 编译器的看法彻底变了。最早做 TorchScript 导出时动不动就遇到“某算子不支持”“控制流被拆成一堆 if”调试一整天最后只能退回 eager 模式。直到 TorchDynamo 出现PyTorch 的图捕获方式才从“用户迁就编译器”翻转为“编译器适配用户代码”。这篇文章我想完整梳理一下这条演进线为什么 TorchScript 不够用TorchDynamo 的捕获原理到底是什么Inductor 和 AOTAutograd 在底层做了哪些事以及落到实际项目里应该怎么调、怎么排查问题。适合正在做模型训练加速、推理服务化或者对 PyTorch 内部机制好奇的工程师。1. 旧时代TorchScript 的教训1.1 它想解决的问题PyTorch 默认的 eager 模式是逐算子执行的。每个算子都要走一遍 Python 层调度、设备分发、kernel launch 这条链路。一条包含几十个算子的前向计算就需要几十次 Python 与 C 之间的跨语言调用。CPU 上还好GPU 上 kernel launch 的延迟会被放大小算子越多越明显。有个很直观的类比eager 模式就像让一群人每次只搬一块砖每个人都要走到工地、放下砖、再走回来编译模式则是先把所有砖按路线排好再派几辆卡车一次性拉完。TorchScript 是 PyTorch 早期给出的答案。它通过torch.jit.trace或torch.jit.script把 Python 代码转换成一种中间表示再经由解释器或底层编译器执行。trace 模式跟踪一次实际输入下的算子调用序列把记录到的计算做成静态图script 模式则直接解析 Python 代码支持if、for这些控制流。设计思路上TorchScript 希望开发者把模型代码写成“可被编译”的形式然后用一套独立于 Python 的运行环境去执行。想法很好但落地时出现了很多问题。1.2 实际项目里踩过的坑我在一个序列模型项目里尝试用torch.jit.script做推理加速模型里有几个带if的动态分支依赖于 tensor 的运行时属性。脚本化之后报了一堆错有的写法必须用torch.jit.isinstance做类型窄化有的需要显式声明Optional类型还有的地方因为索引方式不满足 TorchScript 类型系统被迫改写成masked_select加拼接。最终虽然跑通了但代码面目全非后续维护成本很高。另一个更隐蔽的问题是 trace 模式的“静默错误”。某个模型里有一个依赖输入长度的循环trace 时输入长度为 10循环被完整展开成 10 步。上线后遇到长度为 8 的输入图仍然按 10 步执行最后用切片截断。结果数值还对了调试了很久才发现是循环边界写死导致的。TorchScript 的核心矛盾在于它要求模型代码适配一个新的、受限的 Python 子集。对于常见 CNN 也许尚可但涉及动态 shape、复杂控制流、第三方库算子时就很容易突破边界。很多团队给出的优化结论是“要么重写模型要么放弃编译”这其实是编译器对用户代码的倒逼。1.3 为什么需要换一条路从编译器视角看图捕获是后续一切优化的前置条件。没有完整图就没法做算子融合、内核自动调优、显存规划。TorchScript 思路是“让用户写编译器认识的代码”而 Python 的动态特性、PyTorch 生态里成千上万的自定义算子注定这种约束式方案覆盖不全。TorchDynamo 的出现改变了这个思路。它不要求用户改代码而是在 Python 字节码层面拦截执行过程动态地把可捕获片段组合成子图再把子图交给后端编译。对于捕不到的部分直接回退 eager 执行而不是报错终止。这种“能编多少编多少编不了就算了”的设计让编译对用户代码的侵入性降到了极低也正好是它后来能被设为默认路径的根本原因。2. 核心机制TorchDynamo 是怎么把 Python 变成图的2.1 从字节码层面重新理解执行TorchDynamo 的切入点很特别不解析 Python 源码而是直接操作 Python 字节码。它用PEP 523提供的 frame evaluation hook在每个 Python frame 执行前接管控制权然后分析这个 frame 里将要执行的字节码指令序列。举例来说一个简单函数def forward(x, w, b): y torch.matmul(x, w) z y b return torch.relu(z)Dynamo 会逐条解释字节码。torch.matmul是它认识的算子调用记录下来y b对应BINARY_OP也记录下来torch.relu同样可识别。于是这个 frame 里没有失控的 Python 逻辑可以整体编译成一个子图。最后 Dynamo 把子图交给后端生成一个优化过的执行体替换原来的 Python 函数。如果 frame 里有 Dynamo 不支持的操作——比如调用了某个纯 Python 函数、访问了外部全局对象、使用了依赖运行时值的条件分支——它会在这里截断把前面已经收集到的算子和后面的代码拆开。这个分界点叫做 graph break。Graph break 之前的部分走编译之后的部分继续走 Python 解释器。这个“能截断就截断、不打断程序”的容错设计是 TorchDynamo 最核心的优势之一。它不再追求“百分之百捕获”而是把捕获率做到足够高剩下的交给回退机制兜底。2.2 Guard 机制缓存为什么敢这么狂Dynamo 会把捕获到的子图缓存下来避免每次执行都重新解析。但 Python 对象的状态可能随时变化比如某个 tensor 的 shape、dtype、device 变了或者某个全局配置被改掉了。直接复用旧图就会出错。所以 Dynamo 为每个编译产物维护一组 guard 条件。常见的 guard 有这几类类型 guard检查输入 tensor 的 dtype 是否匹配。形状 guard检查输入 tensor 的 shape 是否与编译时一致。对象属性 guard检查某个 Python 对象的特定属性是否仍保持原值。常量 guard检查被捕获的常量是否仍为同一对象。每次执行时Dynamo 先检查这些 guard。全部通过直接调用缓存的编译结果任何一个不满足就重新捕获、重新编译。这正是动态 shape 场景下重新编译频繁发生的原因——只要 shape 一变旧的 guard 失效就会触发新的编译。所以实际用torch.compile时如果发现第一次调用后仍然有明显的停顿通常不是 bug而是新的 shape 触发了重编译。2.3 为什么它能覆盖绝大多数真实模型TorchDynamo 捕获能力强的另一个原因在于它对“能否编译”的判断粒度足够细。它不是一个函数必须整体可编译才编译而是按操作序列逐段判断。算子被识别到就吸收进图识别不到就在当前点打断后面能吸收的继续吸收。也就是说一个包含大量 Python 控制流、动态逻辑的复杂模型最终可能被切分成多个子图每个子图内部高度优化子图之间由 Python 解释器连接。这个“图破碎但优化仍在”的微切线思路让它在真实业务模型上的图覆盖率可以达到 99% 以上而这正是 TorchScript 无法做到的。3. 图之后的链路AOTAutograd 和 Inductor 做了什么3.1 反向传播也需要一张图Dynamo 把前向算子捕获成子图之后如果想把“全图”优化推得更深反向传播也该一并处理。PyTorch 的反向计算在 eager 模式下由torch.autograd在运行时动态构建计算图对编译器来说看到的是一串零散的算子调用难以整体做融合优化。AOTAutograd 解决的就是这个问题。它在编译期把前向子图对应的反向计算“预先导出”把反向过程也固化成一张静态图。这样一来前向图和反向图可以被当作一个整体看待编译器能同时对两侧做算子融合和内存规划。训练场景里这个机制尤其重要因为反向往往占了一半以上的计算量和显存访问。实际使用中torch.compile会自动套用 AOTAutograd 这套流程不需要手动干预。但对训练性能优化有兴趣的工程师理解这层逻辑有助于判断“为什么显存占用会下降”“为什么编译后的 backward 可以更快”。3.2 Inductor把图变成真正的机器码子图被捕获、反向被导出之后真正把运算压到硬件上的是后端。PyTorch 默认后端 TorchInductor 采用代码生成思路输入一张计算图输出面向特定硬件的优化内核。Inductor 内部设置了两层抽象。上层是一个基于 Python 的中间表示算子被分解为更原语的运算方便做融合分析底层则根据目标硬件输出代码——GPU 上生成 Triton 内核CPU 上生成 C 代码。这里选 Triton 而不是手写 CUDA是为了在表达能力和开发效率之间取得平衡。Triton 本身为深度学习算子设计能描述复杂的 tile 计算和访存模式又不需要开发者手工管理 block 和 thread 的细节。Inductor 最拿手的优化是算子融合。比如“y b再relu”典型实现是读两次显存、写一次显存融合后只需要读一次、写一次。小算子越多的模型收益越明显。对于大矩阵乘法这类访存占比本来就不高的算子融合收益有限这也是为什么某些包含大量矩阵乘法的模型torch.compile的提升幅度不如预期的原因。3.3 训练和推理场景的差异编译优化在训练和推理两个场景里追求的目标是不完全一样的。训练侧重吞吐和显存效率希望满负荷跑、少浪费显存推理往往关注单次延迟要求首 token 延迟低、batch 内的波动小。torch.compile对这两个场景都适用但使用姿势不一样。训练场景一般直接对模型调用torch.compile并把mode设为default减少编译时间、让重编译频率可控推理场景对延迟敏感通常会配合CUDA graph使用同时把 dynamic shape 关掉或明确声明 shape 范围避免运行时触发重编译。这里有个容易踩的坑CUDA graph 会捕获一次执行时的内核序列如果后续输入 shape 变化graph 不能复用需要重新捕获。max-autotune模式下Inductor 会花较长时间搜索最优内核配置推理上线前最好预先跑一边 warmup把编译时间放到部署前完成而不是放到线上首请求。4. 实操torch.compile 的正确打开方式和问题排查4.1 基本用法和参数到底怎么选接入torch.compile很简单一行改代码就够了model torch.compile(model)但这只是默认配置。实际项目里至少需要考虑mode和dynamic两个参数。mode有三个取值default编译时间短、优化适中适合大多数场景reduce-overhead会额外尝试用 CUDA graph 降低启动开销但会占用更多显存max-autotune会花大量时间搜索最优配置适合离线推理或不需要反复重编译的场景。dynamic参数控制是否启用动态 shape 支持。默认情况下Dynamo 会把输入 shape 当作静态值编译一旦 shape 变化就触发重编译。对于 NLP 里常见的变长 batch、多尺度输入的 CV 模型需要设置dynamicTrue或提供 shape 范围否则会出现“第一次调用很慢后面每个新长度又很慢”的现象。很多团队的落地流程是先跑default看收益再针对推理接口切换到reduce-overhead最后在离线批量推理时开启max-autotune。不建议一上来就max-autotune否则编译时间会淹没你验证收益的耐心。4.2 常见报错和排查经验用torch.compile时最常遇到的是 warning 而不是 error类似“Xxx suppressed”或“Graph break due to ...”的输出表示某段代码没办法被捕获Dynamo 回退到了 eager。看见这个先不用慌可以先抓一下图破碎的位置。把环境变量设置成详细日志模式或者直接数一下图破碎的次数定位到具体是哪一行代码导致。还有一类问题是显存不降反升。这通常出现在reduce-overhead或max-autotune模式下因为 CUDA graph 需要占用额外的显存做内存预留。遇到这种情况检查自己的显存余量如果紧张就切回default。另外torch.compile在动态 shape 场景下的显存开销也比静态 shape 高因为每次重编译都可能生成新内核、占用更多内存。数值不一致是另一个需要警惕的点。编译优化的计算顺序可能和 eager 不同带来极小的浮点误差这是正常现象。但如果误差大到影响训练收敛或推理结果就要考虑是不是某个自定义算子被错误地融合了。这时候可以先禁用融合逐个排查。4.3 性能对比的正确姿势验证编译收益不能只看端到端时间的一次跑分就下结论。有两点很容易被忽略warmup 和统计口径。编译后的模型首次或前几次调用包含编译时间不能算进稳定性能。建议先跑 20~50 次 warmup再用后面的平均值作为性能基准。另外要注意端到端时间改善可能被数据加载或预处理淹没。比如你的跑分脚本里数据加载占了 60% 的时间模型计算只占了 40%即使模型计算快了一倍端到端也只是从 1.0 变成 0.8看起来只有 20% 的提升容易误判编译器没效果。正确的做法是先用 profiler 分离数据加载和模型计算的时间再分别对比。我习惯的做法是写一个计时脚本把一次完整迭代拆成“数据加载”“前向计算”“反向传播/推理后处理”三段每一段都单独计时。这样既能看到编译的真实收益也能帮助定位瓶颈到底在哪一层。5. 演进方向从 JIT 到编译基础设施TorchScript 到 TorchDynamo本质上是从“一套独立的 Python 方言”走向“一套嵌入 Python 执行流程的编译调度器”。Dynamo 不试图取代 Python而是把 Python 当作调度层把真正适合编译的算子执行序列提取出来。这个思路对生态兼容的影响是巨大的用户代码不需要重写自定义算子不再成为编译的阻碍。回头去看PyTorch 编译器的演进路径其实反映了深度学习框架对“动态语言 高性能计算”这对矛盾不断妥协又不断突破的过程。TorchScript 代表的是“约束优先”的思路TorchDynamo 代表的是“兼容优先”的思路而 Inductor 的出现说明 PyTorch 团队已经不只是想解决“能不能捕获图”的问题而是要进一步回答“捕获到的图如何编译得足够好”。对我个人来说真正的转折点是做完一次推理服务的优化。之前用 TorchScript 做部署模型代码被改得不像样维护成本极高后来换到torch.compile几乎零改动地获得了明显的延迟收益。那一刻我才意识到编译器的能力边界不再取决于用户愿不愿意重写代码而是取决于编译基础设施本身的工程深度。现在做 PyTorch 项目我几乎默认会把torch.compile纳入性能优化流程但不会再盲目追求全图优化。图破碎并不可怕重要的是知道它在哪、为什么破、影响的算子比例有多大。Dynamo 的价值恰恰在于它把这个问题从“全面失败”变成了“局部可控”。理解这一点比记住任何参数配置都更有用。
阅读完成 · 觉得有帮助？
