PyTorch torch.compile 降低 Guard 开销实战install_free_tensors、guard_filter_fn 与 skip_guard_eval_unsafe 详解【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch每次调用torch.compile编译后的函数Dynamo 都会先执行一组guards守卫再分发到已编译产物。Guards 用来检查编译时所做的前提假设张量 shape 与 dtype、nn.Module属性、全局状态等是否依然成立以保证复用已编译代码是安全的概念背景可参见 Guards。这套检查发生在每一次调用因此对于函数本身计算量相对较小、但 guard 集合较大的场景guard 求值会占据运行时间的可观比例。与它密切相关的是图前字节码pre-graph bytecode——Dynamo 在进入编译图之前运行的一小段字节码负责整理输入例如查找要传入图中的参数和 buffer它同样每次调用都会执行。本文聚焦于降低这类每次调用的固定开销这与降低编译耗时是两个不同的问题如果你的问题是慢编译或反复重编译请参阅 Dealing with Recompilations。警告下文多数选项在设计上就是不安全的它们以牺牲 soundness 为代价换取速度即通过丢弃或跳过 guards 来实现。这些选项假设你的模型代码在 warmup 之后各次调用之间不会改变被守卫的状态如nn.Module属性或全局变量。一旦该假设被打破torch.compile可能静默执行一份过期的编译产物产生错误结果。只有在完全理解每个选项所依赖的假设之后才应启用它们。先测量定位每次调用开销花在了哪里动手优化之前先确认开销是否存在、分布在哪里想看生成了哪些 guards使用 tlparse或设置环境变量TORCH_LOGSguards详见 tlparse / TORCH_TRACE想看guard 求值花了多少时间对编译后函数做 profiling查找TorchDynamo Cache Lookup事件它计时的就是每次调用的 guard 求值from torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CPU]) as prof: opt_mod(x) prof.export_chrome_trace(trace.json) # 在 chrome://tracing 中查看把TorchDynamo Cache Lookup的时间与编译函数的总运行时间对比再决定这个开销是否值得优化以及优化目标应该指向 guards 还是图前字节码。一、用 install_free_tensors 减少图前字节码时间import torch._dynamo with torch._dynamo.config.patch(install_free_tensorsTrue): # 在这里 torch.compile 你的模型 / 调用编译后的函数 ...install_free_tensors把自由张量如参数、buffer安装为图属性graph attributes而不是图输入从而减少图前字节码为把这些张量整理进图所需的开销。作为副作用它也改变了这些参数和 buffer 的 guard 方式——不再需要对每个输入张量逐一匹配因此通常还能带来小幅的 guard 开销改善。但单独看这个改善往往很小因为每次调用的主要成本在于访问被检查的对象而不是 guard 检查本身。从源码可以印证其默认值与来龙去脉Dynamo 配置 中定义为install_free_tensors False # Temporary flag to control the turning of install_free_tensors to True for # export ... install_free_tensors_for_export True也就是说它默认关闭必须像上面那样通过torch._dynamo.config.patch显式开启。该选项最初是作为 export 辅助能力出现的目的是产生数量一致的图输入functional_export.py 甚至在 export 路径上断言install_free_tensors必须为Falseexport 走的是install_free_tensors_for_export。在 variables/builder.py 中config.install_free_tensors直接参与决定某个张量是否以安装方式写入图的属性表这与文档描述的变成图属性而非图输入完全对应。二、用 guard_filter_fn 跳过 nn.Module 的 guardstorch.compile接受guard_filter_fn选项以逐条 guard的粒度决定保留哪些 guards。torch.compiler提供了若干现成的过滤器对 guard 开销影响最大的是跳过nn.Module上的 guardsimport torch opt_mod torch.compile( mod, options{guard_filter_fn: torch.compiler.skip_guard_on_all_nn_modules_unsafe}, )这通常能显著降低 guard 开销因为模型往往有大量被守卫但从不被修改的模块属性。注意它与图前字节码的相互作用跳过这些 guards 之后参数/buffer 指针的缓存会变冷原本从 guards 中省掉的一部分开销会转移到图前字节码里。要看到完整收益需要把本节的过滤器与第一节的install_free_tensorsTrue组合使用——(1) 和 (2) 一起上一个削减 guard 开销另一个防止省下的开销以图前字节码的形式回弹。torch.compiler 提供的全部现成过滤器以下过滤器均在 torch/compiler/init.py 中实现并导出见该文件__all__列表全部带_unsafe后缀、有相同的安全警示。结合源码可以直接读出每个过滤器的判定逻辑skip_guard_on_inbuilt_nn_modules_unsafe— 只跳过内建模块如torch.nn.Linear上的 guards。源码实现是对每条 guard 判断entry.orig_guard.source.is_unspecialized_builtin_nn_module()是则丢弃实现。skip_guard_on_all_nn_modules_unsafe— 跳过所有nn.Module用户自定义 内建上的 guards判定条件是entry.orig_guard.source.is_unspecialized_nn_module()实现。keep_tensor_guards_unsafe— 只保留张量 guards可选是否保留参数 guards。源码逻辑仅当entry.guard_type TENSOR_MATCH时保留其中nn.Parameter类型的守卫默认丢弃传keep_parametersTrue才保留实现。keep_portable_guards_unsafe— 只保留可以跨 Python / 非 Python 环境移植的 guards即全局状态global-state、shape 环境shape-env与非全局张量 guards。源码中保留条件为g.guard_type in (GLOBAL_STATE, SHAPE_ENV) or (g.guard_type TENSOR_MATCH and not g.is_global)实现。skip_guard_on_globals_unsafe— 跳过所有全局变量 guards实现为一行return [not entry.is_global for entry in guard_entries]实现。skip_all_guards_unsafe— 丢弃所有guards移除全部安全保证须极端谨慎使用实现docstring 中明确写有 WARNING: This function will drop all the safety guarantees。三、先尝试 use_recursive_dict_tags_for_guardsimport torch._dynamo torch._dynamo.config.use_recursive_dict_tags_for_guards True该选项通过递归检查 dict tags 来避免运行完整 guard 集合从而加速嵌套nn.Module场景下的 guard 执行。它依赖一套相当复杂的、使用底层 CPython 特性的机制在 OSS issues 中曾引起一些争议因此默认关闭torch/_dynamo/config.py 中use_recursive_dict_tags_for_guards False后续也可能被重新审视。从源码结构看启用后 guard 系统在 finalize 阶段会额外执行 tag 安全根的分析torch/_dynamo/guards.py 的GuardRoot.finalize中仅当config.use_recursive_dict_tags_for_guards为真且通过 justknobs 开关pytorch/compiler:use_recursive_dict_tags_for_guards检查时才会调用find_tag_safe_roots()识别出tag safe node / tag safe root据此用 tag 比较替代全量 guard 重跑。文档建议的尝试顺序是把它放在选项 (1) 和 (2) 之前试——如果对你的模型生效就不需要跳过nn.Moduleguards第二节但仍会受益于install_free_tensors第一节。四、warmup 后用 skip_guard_eval_unsafe 跳过 guard 求值import torch # 1. 预热用足够多样的输入运行编译后模型直到不再发生重编译。 # 2. 然后切换到只运行最小区分 guard 集合的立场。 with torch.compiler.set_stance(skip_guard_eval_unsafeTrue): # 稳态推理 / 训练迭代 ...当你已经把编译模型预热到不再发生重编译的程度后skip_guard_eval_unsafe只运行足以区分你已有各编译产物的最小 guard 集合其余全部跳过。与上面几个选项不同它不能在torch.compile时设置——必须在使用/训练循环里、warmup 完成之后通过 stance 开启。从源码看torch.compiler.set_stance是一个支持函数、上下文管理器和装饰器三种用法的状态切换 APItorch/compiler/init.py其skip_guard_eval_unsafe参数在 docstring 中标注 A flag to run only differentiating guards. CAUTION - This flag is unsafe...。在 torch/_dynamo/eval_frame.py 中stance 被建模为携带skip_guard_eval_unsafe: bool False的DynamoStance状态默认关闭且_set_stance带_dynamo_forbidden标记——这与文档不要在torch.compile区域内调用 set_stance否则会报错的说明一致。如果不再重编译的假设被打破来了一个真正的新输入就有静默产生错误结果的风险这正是名字里unsafe的由来。组合使用推荐的排查与优化路径文档给出的整体实践路径如下先测量用 tlparse /TORCH_LOGSguards看每次调用的时间花在 guards 还是图前字节码上可选地先试use_recursive_dict_tags_for_guardsTrue第三节若无效则install_free_tensorsTrue与guard_filter_fn如skip_guard_on_all_nn_modules_unsafe组合应用——一个削减 guard 开销另一个防止节省重新以图前字节码开销的形式出现稳态服务、warmup 之后考虑set_stance(skip_guard_eval_unsafeTrue)。import torch import torch._dynamo with torch._dynamo.config.patch(install_free_tensorsTrue): opt_mod torch.compile( mod, options{guard_filter_fn: torch.compiler.skip_guard_on_all_nn_modules_unsafe}, ) # 预热若干步确认不再重编译 for x in warmup_inputs: opt_mod(x) # 稳态阶段 with torch.compiler.set_stance(skip_guard_eval_unsafeTrue): for x in serving_inputs: opt_mod(x)小结与延伸阅读本文覆盖的四个手段作用点不同install_free_tensors削减图前字节码并附带小幅 guard 收益guard_filter_fn直接裁剪 guard 集合use_recursive_dict_tags_for_guards用 dict tag 加速 guard 执行skip_guard_eval_unsafe在稳态下只保留区分支产物所需的最少 guards所有_unsafe选项的共同前提是被守卫状态在 warmup 后不再变化违反该假设会导致静默错误结果若要建立对 guards 更完整的心理模型——编译单元图 guard 集合、guard 集合为何庞大以及上述各技术的 profiler 基准实测——原文档推荐进一步阅读 PyTorch 开发者博客《Inside torch.compile Guards》PyTorch 官方博客 devlogs 栏目 2025-06-04 一文相关仓库内资料torch.compiler 模块、Dynamo 配置项、guard 系统实现、stance 状态实现以及同系列的 Dynamo Core Concepts、Observabilitytlparse / TORCH_TRACE 与 Recompilation。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考