PyTorch TorchInductor 溯源追踪（Provenance Tracking）实战指南：从计算图到生成内核的可视化映射

PyTorch TorchInductor 溯源追踪Provenance Tracking实战指南从计算图到生成内核的可视化映射【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读torch.compile会把一个 Python 模型依次经历 pre-grad 图、post-grad 图最终由 TorchInductor或 AOTInductor生成 Triton/C 内核中间经过了大量 pass 的变换普通开发者很难把最终生成的内核与最初的算子对应起来。本文基于 PyTorch 官方文档中的 Provenance Tracking溯源追踪章节结合本仓库源码讲解如何通过tlparse工具将「输入 GraphModule → post-grad 图 → Inductor 生成代码」三段之间的映射关系可视化并逐内核查看其对应的源码与调试句柄debug handle。读完本文你将掌握INDUCTOR_PROVENANCE1环境变量的完整用法、tlparse --inductor-provenance的输出解读以及映射数据在 torch/_inductor/debug.py 与 torch/_inductor/compile_fx.py 中的产生原理。什么是 Inductor 溯源追踪溯源追踪Provenance Tracking是 PyTorch 编译栈中面向调试与性能分析的一项功能服务于 TorchInductor 与 AOTInductor 两条编译路径。其核心能力是可视化输入 GraphModule编译前计算图与优化后生成代码之间的对应关系让你能够追踪原始算子在整个编译过程中被如何变换、融合成最终的内核。工具界面将信息划分为三个面板panel 1输入 GraphModule 中的节点pre-grad 图panel 2post-grad 图梯度计算完成之后、Inductor 调度之前的图中的节点panel 3Inductor 生成的代码wrapper 代码与内核源码。面板中加粗的行表示当前溯源追踪功能已覆盖的节点/内核。目前覆盖的内核类型包括Triton 内核、Ccpp内核与组合内核combo kernels黄色高亮则显示该节点/内核的溯源provenance归属。以下为 TorchInductor 路径下溯源追踪工具的运行效果截图以下为 AOTInductoraot_compile 部署路径下的运行效果使用溯源高亮器Provenance Tracking Highlighter启用并使用溯源追踪功能只需三步安装tlparse、以指定环境变量运行程序、对日志运行tlparse。第一步安装 tlparsetlparse是一个用 Rust 编写的日志解析与可视化工具通过 Cargo 安装cargo install tlparse如果环境中还没有cargoRust 工具链需要先安装 Rust/Cargo参见 Cargo Book 的安装章节本文不再赘述。第二步以溯源追踪模式运行程序在运行 PyTorch 程序时设置两个环境变量TORCH_TRACE~/my_trace_log_dir INDUCTOR_PROVENANCE1 python your_program.pyTORCH_TRACE~/my_trace_log_dir指定 trace 日志输出目录INDUCTOR_PROVENANCE1开启 Inductor 层面的溯源追踪。运行结束后~/my_trace_log_dir下会生成日志文件该日志将被tlparse用于生成溯源追踪高亮器。第三步运行 tlparse 生成可视化对日志文件直接运行tlparse并带上--inductor-provenance标志tlparse log_file_name.log --inductor-provenance运行完成后在tlparse的输出中会出现一个额外的Provenance Tracking溯源追踪区块点击其中的链接即可进入溯源追踪工具界面。使用要点与注意事项不加标志也能看到 JSON 映射即使不加--inductor-provenance标志在tlparse输出的index.html中inductor_provenance_tracking_node_mappings_number.json文件里依然可以看到节点映射的 JSON 数据。直接对日志文件运行 tlparse请直接对 log 文件执行tlparse log_file_name.log --inductor-provenance直接对文件夹运行tlparse parse folder_name --inductor-provenance可能无法正常工作。溯源高亮器依赖的工件artifacts该功能依赖以下由 PyTorch 编译过程产生的文件工件文件名内容before_pre_grad_graph.txtpre-grad输入图文本after_post_grad_graph.txtpost-grad 图文本inductor_aot_wrapper_code.txtAOT wrapper 代码inductor_output_code.txtInductor 最终生成的输出代码inductor_provenance_tracking_node_mappings.json节点映射数据内核 ↔ post-grad ↔ pre-grad以下为 tlparse 输出中 Provenance Tracking 区块入口的示意图溯源追踪的源码实现原理仅仅会使用还不够理解映射数据从何而来能帮助你在数据异常时快速定位问题。溯源追踪的开关与参数集中在 torch/_inductor/config.py 的trace配置中# torch/_inductor/config.py # Save mapping info from inductor generated kernel to post_grad/pre_grad fx nodes # Levels: # 0 - disabled (default) # 1 - normal # 2 - basic # Backward compatibility: # If TORCH_COMPILE_DEBUG1, level is set to at least 1. # If INDUCTOR_PROVENANCE is set, use its integer value. provenance_tracking_level: int int( os.environ.get( INDUCTOR_PROVENANCE, os.environ.get(TORCH_COMPILE_DEBUG, 0) ) )关键点级别语义provenance_tracking_level支持0默认关闭、1normal、2basic三档向后兼容设置TORCH_COMPILE_DEBUG1时级别至少提升为1若同时设置INDUCTOR_PROVENANCE则以其整数值为准。effective_provenance_tracking_level()见 config.py统一了这一逻辑——当TORCH_COMPILE_DEBUG_EXTEND1将内核栈回溯写回 profiler timeline时级别会被强制提升到至少 1。防 OOM 保护provenance_tracking_max_events环境变量TORCH_COMPILE_DEBUG_MAX_EVENTS默认 500000限制了 profiler timeline 后处理阶段处理的 trace 事件数量超过上限即跳过溯源以避免内存溢出。映射数据的采集与序列化实现在 torch/_inductor/debug.py全局状态_inductor_post_to_pre_grad_nodes、_inductor_triton_kernel_to_post_grad_node_info、_inductor_kernel_stack_trace等模块级字典debug.py在编译期间持续累积映射关系reset_provenance_globals()上下文管理器debug.py保证每次编译之间状态被正确重置与恢复。内核级采集set_kernel_post_grad_provenance_tracing()debug.py为每个生成的内核建立与 post-grad 节点的映射对调度节点通过 IR 节点上的origins哪些 FX IR 节点贡献了该融合内核与更精确的origin_node内容与某个 FX 节点输出完全对应收集来源同时收集内核的栈回溯stack traces。该函数还会返回一个自增的整数 debug handle用于区分同一内核的多次调用。JSON 落盘dump_inductor_provenance_info()debug.py在 trace 开启时把合并后的映射写入inductor_provenance_tracking_node_mappings.json并写入version: 2.0字段供 tlparse 识别当前映射格式版本。异常会被signpost_event记录而绝不会影响正常程序执行——这也是整个溯源功能的设计原则。编译入口在 torch/_inductor/compile_fx.py 中before_pre_grad_graphcompile_fx.py 附近与after_post_grad_graph文本被保存inductor_provenance_tracking_node_mappings与inductor_provenance_tracking_kernel_stack_traces两个工件在 compile_fx.py 附近按effective_provenance_tracking_level() ! 0的条件打包进 trace 工件中。输出代码中的溯源信息在 torch/_inductor/output_code.py 中OutputCode类带有inductor_provenance_mapping_str与inductor_provenance_stack_traces_str两个字段output_code.py说明生成代码与映射/栈回溯信息是绑定在一起的。查看每个 Inductor 内核对应的源码启用INDUCTOR_PROVENANCE1之后你还可以在 tlparse 中查看每个 Inductor 内核对应的源码。在 tlparse 输出中找到inductor_provenance_tracking_kernel_stack_traces.json点击其旁边的readable_html链接即可进入内核源码查看页面。理解内核名称中的 debug handle在内核源码视图中内核名称末尾常带有形如:1、:467的后缀例如triton_per_fused_add_1triton_per_fused_add_467这些后缀用于区分同一内核被多次调用的情况官方称之为debug handles调试句柄。在源码实现中正是set_kernel_post_grad_provenance_tracing()通过全局计数器_inductor_kernel_provenance_debug_handle为每个内核追加:{n}后缀见 debug.py从而为每一次内核生成调用建立唯一标识。另外在内核源码的注释中也能找到对应的 debug handle方便你在生成代码与追踪数据之间来回核对典型排查场景溯源追踪在以下场景中尤为实用算子消失了训练/推理脚本中写了某个算子但生成的 CUDA 内核里找不到——通过溯源高亮可以看到该节点是否被常量折叠、被其他内核融合或在某个 pass 中被替换性能归属分析 Inductor 生成的 Triton 内核分别对应计算图的哪些算子从而定位某个“意外变慢”的算子最终被并入了哪个内核AOT 部署对比torch._export.aot_compile导出的 AOTInductor 模型与 JIT 编译路径生成的内核差异可通过三面板映射直观对比。延伸阅读tlparse是 Rust 编写的独立工具本仓库通过INDUCTOR_PROVENANCE等环境变量与其对接相关排查指南见 torch.compiler_troubleshooting。若想深入源码可依次阅读 torch/_inductor/config.py开关与级别、torch/_inductor/debug.py映射采集与落盘、torch/_inductor/compile_fx.py工件生成时机与 torch/_inductor/output_code.py内核代码与溯源信息的绑定。溯源追踪依赖TORCH_TRACE的 trace 日志机制其更多用法可参考 torch.compiler_profiling_torch_compile。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考