PyTorch TorchDynamo 细粒度追踪 API 实战指南torch.compiler.disable / allow_in_graph / disallow_in_graph 精讲【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本篇技术指南以 PyTorch 仓库中 torch.compiler_fine_grain_apis.md 为核心骨架系统讲解 TorchDynamo 提供的细粒度追踪控制 API。当你的模型在torch.compile下整体编译时如果某个局部函数或算子无法被编译器处理你可以用这些 API 精准地跳过、放行或强制打断编译让其余部分仍然获得编译加速。读完本文你将掌握torch.compiler.disable、torch._dynamo.disallow_in_graph、torch.compiler.allow_in_graph、torch._dynamo.graph_break与三个编译状态查询标志is_compiling/is_dynamo_compiling/is_exporting的适用场景、底层实现原理与使用禁忌并能在实际项目中正确选择、组合它们。::: note 在本文中torch.compiler.compile与torch.compile可以互换使用两种写法在你的代码中都会生效见 torch.compiler_fine_grain_apis.md 开篇说明。 :::为什么需要细粒度追踪 APItorch.compile会对整个用户模型执行 TorchDynamo 追踪TorchDynamo 拦截每一个 Python 函数帧frame将它们转换为可编译的图。然而模型代码中总可能出现一小部分无法被torch.compiler处理的逻辑。此时你有两个选择放弃整段编译或者只对出问题的那部分关闭编译器让模型其余部分继续走编译路径。细粒度追踪 API 正是为第二种场景设计的它们允许你精确划定“跳过编译”的代码范围。文档给出的 API 清单如下API作用何时使用torch.compiler.disable在装饰的函数以及递归调用的所有函数上禁用 Dynamo。当模型的一小部分无法被torch.compile处理时用它解除阻塞非常有效。torch._dynamo.disallow_in_graph禁止被标记的算子出现在 TorchDynamo 图中Dynamo 会在此处产生图断裂graph break让该算子以 eager无编译模式运行。适合针对算子而torch.compiler.disable适合装饰函数。当自定义算子如torch.ops.fbgemm.*在torch.compile下出问题时该 API 既可用于调试也可用于解除阻塞。torch.compiler.allow_in_graph被注解的可调用对象原样进入 TorchDynamo 图对 Dynamo 而言相当于一个黑盒。注意 AOT Autograd 仍会穿透它进行追踪因此allow_in_graph只是 Dynamo 层的概念。适用于模型中包含已知的、Dynamo 难以支持特性的部分例如 hooks 或autograd.Function。但每一次allow_in_graph的使用都必须仔细审查无图断裂、无闭包。torch._dynamo.graph_break强制产生一次图断裂断裂前后的代码都会经过 TorchDynamo。对部署场景很少有用——如果你觉得自己需要它大概率真正需要的是disable或disallow_in_graph。torch.compiler.is_compiling指示当前图是否作为torch.compile()或torch.export()的一部分被执行/追踪。在模型代码中做条件分支跳过编译/追踪期间不需要的逻辑。torch.compiler.is_dynamo_compiling指示当前图是否经由 TorchDynamo 追踪。它比torch.compiler.is_compiling()更严格仅在使用 TorchDynamo 时才会为 True。需要区分“编译中”具体由哪个前端触发时使用。torch.compiler.is_exporting指示当前图是否经由 export 追踪。它比torch.compiler.is_compiling()更严格仅在使用torch.export时才为 True。需要区分“追踪”由 export 触发时使用。这些 API 的源码实现集中在 torch/compiler/init.py公共入口与 torch/_dynamo/decorators.pyDynamo 底层实现中下文逐一深入。torch.compiler.disable递归跳过整个调用子树torch.compiler.disable会在被装饰的函数帧及其递归调用的所有函数帧上禁用编译。它的语义建立在“TorchDynamo 会拦截每个 Python 函数帧的执行”这一事实之上。假设你的代码结构如下函数fn调用a_fn和b_fn而a_fn又调用aa_fn与ab_fn。在 eager 模式下这些函数帧原样运行使用torch.compile后TorchDynamo 会逐个拦截这些函数帧下图中的绿色部分设想a_fn是导致torch.compile出问题的函数且它在模型中属于非关键路径。你可以对a_fn使用compiler.disable如上图最右侧所示TorchDynamo 将不再审视从a_fn调用出发的所有帧白色部分表示恢复原始 Python 行为而fn、b_fn等其余部分仍然被编译。最简单的用法是直接装饰出问题的函数torch.compiler.disable def a_fn(x): # 这段代码将以 eager 方式运行不参与编译 return ... torch.compile def fn(x): y a_fn(x) # 此处发生跳过a_fn 及其子调用不进入编译图 z b_fn(x) # b_fn 仍然被 TorchDynamo 编译 return y z如果你不想改动源码也可以使用非装饰器语法手动调用并替换函数引用但官方建议尽量避免这种风格——你必须自行确保原函数的所有调用方都改用修补后的版本否则很容易漏掉某些调用路径。从 torch/_dynamo/decorators.py 的实现可以看到disable的核心参数有两个recursiveTrue默认Dynamo 在被装饰函数帧及其递归调用的函数上被完全跳过recursiveFalseDynamo 只跳过与该函数代码关联的帧但仍然会处理它递归调用的函数帧。这一差异对排查“为什么某个子函数还是被编译了”至关重要。公共入口 torch/compiler/init.py 原样透传这两个参数fn、recursive、reason。其中reason参数会在 Dynamo 尝试追踪被禁用函数时打印出来方便你在日志中定位禁用原因此外源码还保留了_torchdynamo_disable、_torchdynamo_disable_recursive等内部标记供 Dynamo 在追踪时识别被禁用的函数帧。关于recursiveFalse的帧级行为仓库中的另一张示意图做了更细致的刻画被禁用函数a_fn本身恢复 eager 行为而它调用的aa_fn、ab_fn仍可被 TorchDynamo 编译与编译状态标志搭配使用torch.compiler.disable也可以与状态查询标志结合实现“只在编译期跳过”的防御式写法。比如在 test/dynamo/test_decorators.py 的测试中可以看到torch.compiler.is_compiling()被广泛用于条件分支def forward(self, x): if not torch.compiler.is_compiling(): # ...仅在非编译/非追踪期执行的逻辑例如日志、校验... # ...其余参与编译的逻辑...torch._dynamo.disallow_in_graph对单个算子强制图断裂torch._dynamo.disallow_in_graph与disable的关键区别在于它禁止的是算子operator而不是函数。它把指定算子从 TorchDynamo 提取的图中排除Dynamo 会在该算子处产生图断裂并让这个算子回落到 PyTorch eager 模式执行。典型场景是模型可以正常编译出图但下游编译器在某处失败——例如某个算子缺少 meta kernel或者某个算子的 Autograd dispatch key 设置不正确。此时你可以把这个算子标记为disallow_in_graphimport torch._dynamo # 禁止 torch.sub 出现在 Dynamo 图中 torch._dynamo.disallow_in_graph(torch.sub) torch._dynamo.optimize(...) def fn(a): x torch.add(x, 1) x torch.sub(x, 1) # 此处产生图断裂以 eager 模式运行 x torch.add(x, 1) return x如上例源码 torch/_dynamo/decorators.py 的 docstring 中给出执行fn时会围绕torch.sub断成两个图每个图各含一个torch.add()。这里有一个非常容易踩的坑你找到的必须是 Dynamo 层的算子而不是 ATen 层算子。详见后文「Limitations」小节。::: warningtorch._dynamo.disallow_in_graph是一个全局标志。如果你在对比不同的后端编译器切换编译器时需要为被禁止的算子重新调用allow_in_graph以撤销禁止否则会影响新编译器的行为。 :::从实现上看disallow_in_graph经由_disallow_in_graph_helper(throw_if_not_allowedTrue)注册到 Dynamo 的追踪规则trace rules中并在追踪到该算子时强制触发图断裂测试用例 test/dynamo/test_decorators.py 中的test_disallow_in_graph演示了disallow_in_graph(torch.sub)后再用allow_in_graph(torch.sub)恢复的完整流程。同一文件中还有一个forbid_in_graph它断言某函数在追踪过程中不得出现否则报错如果你想要的是图断裂而不是硬报错应该使用disallow_in_graph。torch.compiler.allow_in_graph把函数当作黑盒放入图torch.compiler.allow_in_graph用于这样的场景相关函数帧包含已知的、TorchDynamo 难以支持的特性例如 hooks、autograd.Function而你有信心下游组件如 AOT Autograd可以安全地穿透该函数完成追踪。当函数被allow_in_graph装饰后TorchDynamo 把它当作黑盒原样放进生成的图中torch.compiler.allow_in_graph(my_custom_function) torch.compile(...) def fn(x): x torch.add(x, 1) x my_custom_function(x) # 原样进入图Dynamo 不再符号化地审视其内部 x torch.add(x, 1) return x执行后最终会捕获一个包含my_custom_function()的单一图。::: warningallow_in_graph会在被装饰函数上完全跳过 TorchDynamo从而省略掉 TorchDynamo 的所有安全检查包括图断裂处理、闭包处理等。请谨慎使用。PyTorch 下游组件如 AOT Autograd依赖 TorchDynamo 处理复杂的 Python 特性而allow_in_graph绕过了 TorchDynamo使用不当可能导致正确性soundness问题且极难排查。 :::关于这条警告torch/compiler/init.py 中的实现给出了更精确的定位allow_in_graph是“前端Dynamo不追踪但后端AOTAutogradInductor仍会穿透追踪”。文档中的对比表格可以总结为机制前端Dynamo后端AOTAutogradInductor无装饰器追踪函数内部追踪函数内部allow_in_graph视为不透明可调用对象追踪函数内部custom op自定义算子视为不透明可调用对象视为不透明可调用对象因此如果目标是“在整个torch.compile栈中都把某个 Python 函数当作黑盒”官方推荐的做法是创建自定义算子custom operator而不是allow_in_graph——这正是 torch/compiler/init.py 中明确的建议。allow_in_graph对被装饰函数有严格要求不满足则行为未定义见 torch/compiler/init.py函数的输入必须是 FX 图中可被 Proxy 化的类型Tensor/int/bool/float/None/List[Tensor?]/List[int?]/List[float?]/Tuple[Tensor?, ...]/Tuple[int?, ...]/Tuple[float?, ...]/torch.dtype/torch.device函数的输出也必须是上述可 Proxy 化类型函数内部使用的所有 Tensor 都必须直接作为参数传入而不能是捕获的外部变量。底层实现torch/_dynamo/decorators.py会把函数 id 注册到trace_rules._allowed_callable_ids从而让 Dynamo 的调用规则lookup_callable将其识别为TorchInGraphFunctionVariable直接放进图中。仓库测试 test/dynamo/test_autograd_function.py 的test_allow_in_graph、test/dynamo/test_decorators.py 的test_allow_in_graph与test_allow_in_graph_no_id_reuse防止 id 复用导致的微妙 bug都验证了这一机制。另外注意一个细节allow_in_graph接受单个可调用对象也接受 list/tuple会递归地对每个元素应用并返回新的 list/tuple。如果你只是需要“在图中放一个不透明的调用”但又想要更宽松的输入约束可以关注文档中提到的nonstrict_trace()见 torch/compiler/init.py它支持用户自定义类、nn.Module参数与全局捕获 Tensor且同样支持反向传播。torch._dynamo.graph_break显式插入图断裂点torch._dynamo.graph_break强制在代码中的指定位置产生一次图断裂断裂之前的代码和断裂之后的代码都会经过 TorchDynamo各自生成图。它的典型用法是在代码中直接调用import torch._dynamo def fn(x): x torch.relu(x) torch._dynamo.graph_break(这里显式断图) # 强制图断裂 x torch.sigmoid(x) return x其实现torch/_dynamo/decorators.py同样是基于_disallow_in_graph_helper(throw_if_not_allowedFalse)但语义上它只是一个“强制图断裂”的标注点与disallow_in_graph针对具体算子不同。官方对它的态度非常明确对部署场景很少有用。如果你觉得需要手动插入图断裂大概率真正需要的是disable跳过整段函数或disallow_in_graph跳过某个算子。在同族工具里还有skip_frame与step_unsupported前者强制跳过整个帧后者在“已追踪的 FX 图先编译、然后跳过帧内剩余代码”这种特殊调试场景使用——它们主要用于内部调试。编译状态查询is_compiling / is_dynamo_compiling / is_exporting与前面几个“控制追踪行为”的 API 不同这一组 API 用于在模型代码内部感知当前是否处于编译/追踪状态从而有条件地跳过只在 eager 期需要的逻辑。它们的实现与标志位集中在 torch/compiler/init.py标志语义实现要点torch.compiler.is_compiling()图是否作为torch.compile()或torch.export()的一部分被执行/追踪读取_is_compiling_flagtorch.jit.is_scripting()时返回False见 torch/compiler/init.py。torch.compiler.is_dynamo_compiling()图是否经由 TorchDynamo 追踪比is_compiling更严格仅 TorchDynamo 参与时为 True见 torch/compiler/init.py。torch.compiler.is_exporting()图是否经由torch.export追踪比is_compiling更严格读取_is_exporting_flag见 torch/compiler/init.py。典型用法是条件分支def forward(self, x): if not torch.compiler.is_compiling(): # ...编译/追踪期不需要的逻辑如耗时校验、日志... # ...其余需要参与编译的逻辑...需要指出的是_is_compiling_flag由_compile_session_context()在每次torch.compile会话期间置位旧的torch._dynamo.external_utils.is_compiling()与torch._utils.is_compiling()已被标记为弃用并转发到torch.compiler.is_compiling()见 torch/_dynamo/external_utils.py新代码应直接使用torch.compiler命名空间下的这三个标志。仓库中的分布式测试 test/distributed/test_inductor_collectives.py 也展示了torch.compiler.is_dynamo_compiling()的实际应用场景。Limitations这些 API 的视野边界所有上述 API 都作用在TorchDynamo 层因此它们只能看到 TorchDynamo 能看到的范围。这会导致一些容易困惑的场景。最典型的例子torch._dynamo.disallow_in_graph对 ATen 层算子不生效因为 ATen 算子是 AOT Autograd 可见的。例如# 下面的写法不生效 torch._dynamo.disallow_in_graph(torch.ops.aten.add)原因在于 TorchDynamo 提取图时torch.ops.aten.add这类底层算子位于 Dynamo 的“视野之外”AOT Autograd 的职责范围disallow_in_graph无法对它们触发图断裂。这正是官方文档提醒“你必须找到对应的 Dynamo 层算子而不是 ATen 层算子”的原因。同样地allow_in_graph也只是 Dynamo 层的概念——它让 Dynamo 不再穿透但 AOT Autograd 仍然会继续追踪这一点在做跨组件Dynamo → AOT Autograd → Inductor的正确性分析时必须牢记。实战决策速查根据本指南的内容在模型局部代码无法被torch.compile处理时可以按如下顺序决策问题在整段 Python 函数含其子调用→ 用torch.compiler.disable装饰该函数只想跳过函数本身、保留子函数编译时用disable(fn, recursiveFalse)问题在某个具体算子如torch.ops.fbgemm.*等自定义算子→ 用torch._dynamo.disallow_in_graph(op)并确认该算子属于 Dynamo 可见层注意它是全局标志切换后端时要记得allow_in_graph恢复函数内部有 hooks / autograd.Function 等 Dynamo 难处理特性且确定 AOT Autograd 能安全穿透→ 用torch.compiler.allow_in_graph但必须逐处审查输入输出类型与闭包使用防止正确性风险想在模型代码里感知编译状态→ 用torch.compiler.is_compiling()/is_dynamo_compiling()/is_exporting()做条件分支torch._dynamo.graph_break→ 仅在极少数调试/特殊场景使用部署代码通常不应依赖它。最后再次强调这些 API 的语义边界Dynamo 层 vs AOT Autograd 层是使用时最大的隐患务必结合本文「Limitations」一节与对应源码torch/_dynamo/decorators.py、torch/compiler/init.py理解后再落地到项目中。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考