TorchInductor 和 AOTInductor 溯源跟踪#

创建时间：2025 年 05 月 09 日 | 最后更新时间：2025 年 05 月 23 日

警告

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

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

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

粗体行表示当前溯源跟踪功能所覆盖的节点/内核。我们目前涵盖 triton 内核、cpp 内核和 combo 内核。黄色高亮显示了节点/内核的溯源。

使用溯源跟踪高亮器#

请按照以下步骤在您的 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 输出中看到一个额外的“Provenance Tracking”部分。点击链接即可访问溯源跟踪工具。有关演示，请参阅：pytorch/tlparse#93

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