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 編寫的工具。