PyTorch TorchDynamo 工作原理详解：从字节码捕获、Guards 机制到缓存工件排查

PyTorch TorchDynamo 工作原理详解从字节码捕获、Guards 机制到缓存工件排查【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchTorchDynamo简称 Dynamo是 PyTorch 内置的 Python 级即时编译JIT编译器它通过 CPython 的帧求值 APIPEP 523 并结合当前仓库源码为你讲透 Dynamo 的 Guards 机制、字节码改写原理、TORCH_LOGS调试手段以及如何通过_debug_get_cache_entry_list等 API 检视编译产物帮助你真正理解torch.compile在底层做了什么。阅读本文前建议先通读 Torch Compiler 总览了解torch.compile的整体组成。一图看懂有 torch.compile 与没有 torch.compile 的区别TorchDynamo 的目标是让未经修改的 PyTorch 程序运行得更快它不对用户代码做任何 API 层面的侵入而是从执行机制层面下手。下图对比了默认 Python 执行流程与 Dynamo 介入后的执行流程左侧是 PyTorch 的默认Eager行为用户函数经PyFrameObject/PyCodeObject直接交给_PyEval_EvalFrameDefault()逐条解释执行没有任何优化。右侧是 Dynamo 的介入路径它在帧求值时对字节码做动态分析与转换通过 Guards 检查执行环境、把 Torch 相关操作提取成 FX Graph由用户自定义的后端编译器生成编译函数Compiled Function并利用缓存避免重复编译最终仍然通过_PyEval_EvalFrameDefault()执行改写后的字节码。Dynamo 在「Python 解释执行的灵活性」与「编译后端的性能」之间取得平衡——它能处理的算子进图编译不能处理的留在 Python 解释器中执行两者协同工作。Dynamo 是什么字节码级别的图捕获器TorchDynamo 的核心定位是基于字节码分析的图捕获器它挂钩 CPython 的帧求值 APIPEP 523 定义的机制在 Python 字节码真正执行之前动态修改它它把字节码改写成提取 PyTorch 算子序列的形式生成 FX Graph图结构表示这个 FX Graph 被交给一个可定制的后端如 TorchInductor、aot_eager、用户自定义编译器编译加速。因此Dynamo 天然支持「混合执行」图中可编译的部分走后端编译图外无法捕获的 Python 逻辑仍由解释器执行兼顾可用性与性能。Dynamo 让尝试不同编译后端变得异常简单——只需一行装饰器。torch._dynamo.optimize()是底层 API而torch.compile()是对它的便捷封装这也是日常使用中接触最多的入口。当前仓库中torch._dynamo包torch/_dynamo即为此机制的全部实现其中 decorators.py 定义了optimize/disable/run等装饰器eval_frame.py 承载了帧求值挂钩与缓存管理核心逻辑。Dynamo InternalsGuards 机制Dynamo 以即时JIT方式工作会基于动态属性对计算图做特化specialize。所谓 Guards守卫就是一组「该特化图成立的前提条件」——只有所有条件都满足之前编译好的图才能复用一旦某个条件被破坏Dynamo 就会重新捕获图并重新编译。最小示例自定义后端编译 toy_example下面是文档中的经典示例用torch.compile(backend...)装饰一个函数其中my_compiler是一个最简单的自定义后端——它接收 FXGraphModule和示例输入打印图结构然后直接返回可调用的gm.forward即让编译产物退化为原始 Python 图方便观察from typing import List import torch def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]): print(my_compiler() called with FX graph:) gm.graph.print_tabular() return gm.forward # return a python callable torch.compile(backendmy_compiler) def toy_example(a, b): x a / (torch.abs(a) 1) if b.sum() 0: b b * -1 return x * b for _ in range(100): toy_example(torch.randn(10), torch.randn(10))注意backend回调的签名与仓库中后端协议一致——接收(gm, example_inputs)返回一个 Python 可调用对象作为编译结果见 torch/_dynamo/backends 中各内置后端的实现模式。Guards 长什么样上面示例第一次捕获的图会附带如下 GuardsGUARDS: hasattr(L[a], _dynamo_dynamic_indices) False hasattr(L[b], _dynamo_dynamic_indices) False utils_device.CURRENT_DEVICE None ___skip_backend_check() or ___current_backend() ___lookup_backend(140355900538256) check_tensor(L[a], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, deviceNone, requires_gradFalse, size[10], stride[1]) check_tensor(L[b], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, deviceNone, requires_gradFalse, size[10], stride[1])其中L是当前帧局部变量的字典L[a]、L[b]即函数入参。任何一条 Guard 失败都意味着当前特化图不再安全Dynamo 会触发重新捕获与重新编译。check_tensor最关键的 Guard上面输出里最有价值的 Guard 是check_tensor它检查torch.Tensor的如下属性Tensor 的Python 类如是否 Tensor 子类 / 张量子类dtype数据类型device设备requires_grad是否需要梯度dispatch_key分派键已应用线程局部的 include/exclude 过滤ndim维度数sizes形状strides步长关于这些条件的含义完全特化模式full specialization允许后端编译器假定整张图完全静态——遗憾的是大多数后端如 Inductor都依赖这一点。当算子返回动态形状dynamic shapes而当前未开启动态形状模式时就会触发一次graph break图断裂执行流退回 Python 解释器动态形状相关的处理见 torch.compiler_dynamic_shapes。在仓库实现层面Guards 的生成与检查逻辑位于 torch/_dynamo/guards.py每个叶子 Guard 通过verbose_code_parts()生成可读的条件片段populate_code_parts_for_debugging()会收集全部片段并填充到code_parts列表供调试打印最终由 guard 函数在每次调用时按序求值。Dynamo 到底在做什么TORCH_LOGS 与字节码反编译想看清 Dynamo 的运作细节最直接的方法是开启日志。在运行程序前设置环境变量TORCH_LOGSdynamo,guards,bytecode这会打印出 Dynamo 捕获图、生成 Guards、改写字节码的完整过程。日志体系由 torch/_dynamo/logging.py 实现与torch._logging的全局日志配置见 torch/_logging打通。用 depyf 把字节码还原成源码如果你不熟悉 Python 字节码可以安装反编译钩子把字节码还原成可读的 Python 源码。文档推荐的社区工具是depyf安装与启用方式pip install depyfimport depyf depyf.install()这会产生信息量大但略显刷屏的打印输出。剖析 toy_example 的捕获产物以toy_example的第一个图为对象TORCH_LOGS输出大致分为四段FX Graph → 原始字节码 → Dynamo 改写后的字节码 → 反编译源码最后是上文提到的 Guards。首先是 FX 图可以看到a / (abs(a) 1)被拆成abs_1、add、truediv三个算子节点b.sum() 0对应sum_1、lt两个节点__compiled_fn_0 eval_with_key.1 opcode name target args kwargs ------------- ------- ------------------------------------------------------ ---------------- -------- placeholder a a () {} placeholder b b () {} call_function abs_1 built-in method abs of type object at 0x7f9ca082f8a0 (a,) {} call_function add built-in function add (abs_1, 1) {} call_function truediv built-in function truediv (a, add) {} call_method sum_1 sum (b,) {} call_function lt built-in function lt (sum_1, 0) {} output output output ((truediv, lt),) {}接着是函数的原始字节码这是解释器眼中原始的toy_exampleORIGINAL BYTECODE toy_example example.py line 12 14 0 LOAD_FAST 0 (a) 2 LOAD_GLOBAL 0 (torch) 4 LOAD_METHOD 1 (abs) 6 LOAD_FAST 0 (a) 8 CALL_METHOD 1 10 LOAD_CONST 1 (1) 12 BINARY_ADD 14 BINARY_TRUE_DIVIDE 16 STORE_FAST 2 (x) 15 18 LOAD_FAST 1 (b) 20 LOAD_METHOD 2 (sum) 22 CALL_METHOD 0 24 LOAD_CONST 2 (0) 26 COMPARE_OP 0 () 28 POP_JUMP_IF_FALSE 19 (to 38) 16 30 LOAD_FAST 1 (b) 32 LOAD_CONST 3 (-1) 34 BINARY_MULTIPLY 36 STORE_FAST 1 (b) 17 38 LOAD_FAST 2 (x) 40 LOAD_FAST 1 (b) 42 BINARY_MULTIPLY 44 RETURN_VALUE以及 Dynamo改写后的字节码——注意原始函数体已被替换为对编译产物__compiled_fn_0和两个 continuation 函数__resume_at_30_1/__resume_at_38_2的调用MODIFIED BYTECODE toy_example example.py line 12 12 0 LOAD_GLOBAL 3 (__compiled_fn_0) 2 LOAD_FAST 0 (a) 4 LOAD_FAST 1 (b) 6 CALL_FUNCTION 2 8 UNPACK_SEQUENCE 2 10 STORE_FAST 2 (x) 12 POP_JUMP_IF_FALSE 12 (to 24) 14 LOAD_GLOBAL 4 (__resume_at_30_1) 16 LOAD_FAST 1 (b) 18 LOAD_FAST 2 (x) 20 CALL_FUNCTION 2 22 RETURN_VALUE 24 LOAD_GLOBAL 5 (__resume_at_38_2) 26 LOAD_FAST 1 (b) 28 LOAD_FAST 2 (x) 30 CALL_FUNCTION 2 32 RETURN_VALUE最后是 depyf 反编译出的「可能的源码」与改写后的字节码一一对应def toy_example(a, b): __temp_1 __compiled_fn_0(a, b) x __temp_1[0] if __temp_1[1]: return __resume_at_30_1(b, x) return __resume_at_38_2(b, x)graph break 与 resume_at 机制上述输出清晰展示了 Dynamo 的图断裂处理__compiled_fn_0是my_compiler()的返回值即编译后的图__resume_at_30_1和__resume_at_38_2是 Dynamo 生成的continuation续行函数分别在字节码偏移 30 和 38 处即if分支内部和函数结尾接续执行。每个 continuation 函数的形式为__resume_at_offset: ... restore stack state if needed ... JUMP_ABSOLUTE offset into toy_example ... original bytecode of toy_example ...其工作原理是通过生成resume_at函数Dynamo 强制把函数剩余部分放到一个新的 Python 帧中执行当执行到达该点并第一次命中时新帧会递归触发 Dynamo 重新开始图捕获。这就是「图断裂后自动重启捕获」的底层机制相关实现见 torch/_dynamo/resume_execution.py。如何检视 Dynamo 生成的工件从函数code中取出缓存条目Dynamo 会把编译产物缓存到函数对象的__code__上。仓库为此提供了调试 APItorch._dynamo.eval_frame._debug_get_cache_entry_list它位于 torch/_dynamo/eval_frame.py实现是「给定 code 对象或可调用对象取出存储在其上的缓存条目列表」。from torch._dynamo.eval_frame import _debug_get_cache_entry_list, innermost_fn cache_entries _debug_get_cache_entry_list(innermost_fn(toy_example)) cache_entry cache_entries[0] guard, code cache_entry.check_fn, cache_entry.code # the guard takes the local variables of an input frame, # and tells whether a re-compilation should be triggered. import dis dis.dis(guard) dis.dis(code)其中innermost_fn的作用是在多层_TorchDynamoContext嵌套调用的情况下找到最内层的函数。因为 Dynamo 的缓存挂在fn.__code__上必须定位到最内层函数才能正确取到缓存定义同样在 eval_frame.py。一个编译后的函数可以有多个缓存条目每个缓存条目由一个检查 Guards 的生成函数check_fn和一个types.CodeTypecode组成——后者是 Guards 全部满足时实际执行的代码。如果你熟悉 Python 字节码直接dis上述输出即可理解。如果你不熟悉还有更友好的方式。直接读取 Guard 条件guard.code_parts对于 guard 函数无需解析字节码可以直接访问它的守卫条件列表for code_part in guard.code_parts: print(code_part)输出___guarded_code.valid ___check_global_state() hasattr(L[a], _dynamo_dynamic_indices) False hasattr(L[b], _dynamo_dynamic_indices) False utils_device.CURRENT_DEVICE None ___skip_backend_check() or ___current_backend() ___lookup_backend(140215810860528) ___check_tensors(L[a], L[b], tensor_check_namestensor_check_names)只有当所有条件都满足时guard 函数才返回 true随之执行编译代码。从源码看code_parts是 guards.py 中GuardManager的成员self.code_parts: list[str] []在populate_code_parts_for_debugging()中由每个叶子 Guard 的verbose_code_parts()生成并填充。反编译编译代码对于编译代码无法直接拿到其源码只能反编译from depyf import decompile print(decompile(code))输出def toy_example(a, b): __temp_1 __compiled_fn_0(a, b) x __temp_1[0] if __temp_1[1]: return __resume_at_30_1(b, x) return __resume_at_38_2(b, x)代码中引用的三类名字反编译代码里出现的名字可以归为三类编译函数存放在原函数toy_example所在模块的全局命名空间中如__compiled_fn_0、__resume_at_30_1、__resume_at_38_2。用于检查 Guards 的闭包变量通过guard.__code__.co_freevars获取名字、guard.__closure__获取值如___guarded_code、___is_grad_enabled、___are_deterministic_algorithms_enabled、___is_torch_function_enabled、utils_device、___check_tensors、tensor_check_names。guard 函数的参数L一个把toy_example参数名映射到值的字典仅在函数被调用时可用此时帧求值 API 介入。结构形如{a: value_a, b: value_b}因此代码中用L[a]指代输入变量a。可以看到编译后的toy_example代码中保留了图断裂点——必须由 Python 解释器来选择接下来执行哪个子图。自定义后端下各子图的源码因为我们传入的my_compiler只是一个简单的 Python 后端所以子图__compiled_fn_0、__resume_at_30_1、__resume_at_38_2仍然都是 Python 代码可以逐一检视忽略函数名只看签名与函数体print(source code of __compiled_fn_0:) print(innermost_fn(__compiled_fn_0).__self__.code) print( * 60) print(source code of __resume_at_30_1:) print(decompile(__resume_at_30_1)) print( * 60) print(source code of __resume_at_38_2:) print(decompile(__resume_at_38_2))输出source code of __compiled_fn_0: def forward(self, L_a_ : torch.Tensor, L_b_ : torch.Tensor): l_a_ L_a_ l_b_ L_b_ abs_1 torch.abs(l_a_) add abs_1 1; abs_1 None truediv l_a_ / add; l_a_ add None sum_1 l_b_.sum(); l_b_ None lt sum_1 0; sum_1 None return (truediv, lt) # To see more debug info, please use graph_module.print_readable() source code of __resume_at_30_1: def resume in toy_example(b, x): b b * -1 return x * b source code of __resume_at_38_2: def resume in toy_example(b, x): return x * b注意一旦换成内置后端如inductor这些子图代码在 GPU 上就是编译后的CUDA kernel在 CPU 上则是生成的C 代码而不再是 Python。概念等价编译后代码的完整形态把上述机制浓缩成伪代码编译后的函数概念上等价于def compiled_example(a, b): L {a: a, b: b} for guard, code in get_cache_entries(): if guard(L): return code(a, b) recompile_and_add_another_cache_entry()即先查缓存、逐个过 Guard命中则直接执行已编译代码全部未命中则重新编译并新增一个缓存条目。这也是 Dynamo 缓存机制的运行模型缓存条目相关类型定义见 torch/_dynamo/types.py。下面的流程图展示了torch.compile对用户代码的完整变换先从用户函数中提取计算图红色块进__compiled_fn_0蓝色分支进__resume_at_30_1黄色返回逻辑进__resume_at_38_2交给后端编译优化再由 Guard 函数校验环境最终组装成一个与原函数功能等价、但计算速度更优的新函数延伸阅读想深入了解 Dynamo 的字节码级实现细节参见 TorchDynamo 深度解析torch.compile的自定义后端协议与更多后端示例见 自定义后端动态形状dynamic shapes模式如何影响 Guards 与 graph break见 动态形状编译后的 FX 图与后端产物CUDA kernel / C 代码的检视方法见 使用 torch.compile 性能剖析。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考