TorchInductor 和 AOTInductor 的来源跟踪#

创建日期: 2025年5月9日 | 最后更新日期: 2025年5月23日

警告

此功能是一个处于积极开发中的原型，未来版本中可能会有重大更改。该工具目前的兼容性仅限于最新的 PyTorch nightly 版本。

本节介绍了如何在 tlparse 中使用 TorchInductor 和 AOTInductor 的来源跟踪功能。来源跟踪可帮助您可视化输入 GraphModule 到 (AOT)Inductor 以及生成的优化代码之间的关系。此功能允许您跟踪原始操作在编译期间如何被转换。

下面显示了来源跟踪工具的一些示例屏幕截图。该工具可视化了输入图的节点（面板 1）、渐变后图（面板 2）以及 Inductor 生成的代码（面板 3）之间的映射关系。

加粗的行表示当前来源跟踪功能覆盖的节点/内核。我们目前涵盖 Triton 内核、cpp 内核和组合内核。黄色高亮显示了节点/内核的来源。

使用来源跟踪高亮器#

请按照以下步骤在您的 PyTorch 项目中启用和使用来源跟踪：

通过 cargo install tlparse 安装 tlparse。如果您没有 cargo，请参阅 The Cargo Book 了解安装说明。
使用必要的标志运行您的程序
```
TORCH_TRACE=~/my_trace_log_dir TORCH_LOGS="+inductor" TORCH_COMPILE_DEBUG=1 python your_program.py
```
这将在 ~/my_trace_log_dir 中生成一个日志文件。该日志文件将由 tlparse 用于生成来源跟踪高亮器。
使用 --inductor-provenance 标志在日志上运行 tlparse。例如：
```
tlparse log_file_name.log --inductor-provenance
```
- 即使您不添加 --inductor-provenance 标志，您也应该能够在 index.html tlparse 输出中的 inductor_provenance_tracking_node_mappings_<number>.json 文件中以 json 格式看到映射。
- 直接在日志文件上运行 tlpare。如果您运行“tlparse parse <folder_name> –inductor-provenance”，它可能无法正常工作。
- 来源跟踪高亮器使用的 tlparse 伪像是：
  - before_pre_grad_graph.txt
  - after_post_grad_graph.txt
  - inductor_aot_wrapper_code.txt
  - inductor_output_code.txt
  - inductor_provenance_tracking_node_mappings.json

运行 tlparse <file_name> --inductor-provenance 后，您应该会在 tlparse 输出中看到一个额外的“来源跟踪”部分。点击链接访问来源跟踪工具。演示请参见：pytorch/tlparse#93

tlparse 是一个用 Rust 编写的工具。