TorchInductor GPU 效能分析

本節列出了一些有用的命令和工作流程，可以幫助您深入瞭解模型在 TorchInductor 中的效能。當模型執行速度不如預期時，您可能需要檢查模型的各個核心。通常，佔用 GPU 時間最多的核心是最值得關注的。之後，您可能還想直接執行單個核心並檢查其效能。PyTorch 提供了涵蓋上述所有內容的工具。

相關環境變數

您可以在分析中使用以下環境變數

• TORCHINDUCTOR_UNIQUE_KERNEL_NAMES

  • 預設情況下，TorchInductor 將 Triton 核心命名為 ‘triton\_’。啟用此環境變數後，inductor 會在 trace 中生成更有意義的核心名稱，例如 triton_poi_fused_cat_155，其中包含核心類別（poi 表示逐點操作）和原始 ATen 運算元。預設停用此配置是為了提高編譯快取命中率。

• TORCHINDUCTOR_BENCHMARK_KERNEL

  • 啟用此選項將使 inductor 程式碼生成工具對單個 triton 核心進行基準測試。

• TORCHINDUCTOR_MAX_AUTOTUNE

  • Inductor 自動調優器將對更多 triton.Configs 進行基準測試，並選擇效能最佳的配置。這會增加編譯時間，以期提升效能。

分解模型的 GPU 時間

以下是將模型執行時間分解到單個核心的步驟。我們以 mixnet_l 為例。

1. 執行模型的基準測試指令碼

   TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 TORCHINDUCTOR_BENCHMARK_KERNEL=1
   python -u benchmarks/dynamo/timm_models.py –backend inductor –amp
   –performance –dashboard –only mixnet_l –disable-cudagraphs –training
   

   注意

   該工具依賴核心名稱來決定其類別。啟用 TORCHINDUCTOR_UNIQUE_KERNEL_NAMES 對此至關重要。

2. 在輸出日誌中，查詢以下行

   **Compiled module path:
   /tmp/torchinductor_shunting/qz/cqz7hvhood7y3psp7fy6msjxsxyli7qiwiybizdwtjw6ffyq5wwd.py**
   

每個已編譯模組對應一行。如果沒有額外的圖中斷，我們將在日誌中看到 2 行此類資訊，一行對應前向圖，一行對應後向圖。

對於我們的示例命令，我們分別獲得了前向圖和後向圖的以下已編譯模組

• https://gist.github.com/shunting314/c2a4d8a28b00fcb5586d0e9d9bf77f9f

• https://gist.github.com/shunting314/48efc83b12ec3ead950052e4a0220b10

3. 現在我們可以深入瞭解每個已編譯模組的效能。讓我們選擇前向圖進行說明。為了方便，我將其命名為 fwd.py。使用 -p 引數直接執行它

   **> python fwd.py -p**
   

請在此示例 gist 中檢視完整的輸出日誌。

在輸出中，您會注意到以下內容

• 我們為效能分析寫入了一個 Chrome trace 檔案，以便我們可以載入 trace 並與之互動。在日誌中，查詢如下行以找到 trace 檔案的路徑。

效能分析的 Chrome trace 已寫入 /tmp/compiled_module_profile.json

將 trace 載入到 Chrome 中（在 Chrome 瀏覽器中訪問 chrome://tracing 並按照 UI 提示載入檔案）將顯示如下 UI

您可以放大和縮小以檢查效能分析結果。

• 我們透過如下日誌行報告 GPU 時間佔 wall time 的百分比

  GPU 忙碌時間百分比: 102.88%

  有時您可能會看到大於 100% 的值。原因是 PyTorch 在啟用效能分析時使用核心執行時間，而在停用效能分析時使用 wall time。效能分析可能會稍微扭曲核心執行時間。但總的來說，這應該問題不大。

  如果我們使用較小的批次大小執行像 densenet121 這樣的模型，我們會看到 GPU 忙碌時間百分比很低

  (Forward graph) Percent of time when GPU is busy: 32.69%
  

  這意味著模型有很多 CPU 開銷。這與啟用 cudagraphs 可大幅提高 densenet121 效能的事實一致。

• 我們可以將 GPU 時間分解到不同類別的核心。在 mixnet_l 示例中，我們看到

  • 逐點核心佔用 28.58%

  • 歸約核心佔用 13.85%

  • 持久歸約核心佔用 3.89%

  • 其餘是用於 mm/conv 的 cutlass/cudnn 核心，佔用 56.57%

  此資訊可在每種核心類別的報告的彙總行（最後一行）中找到。

• 我們還可以放大檢視特定類別的核心。例如，讓我們檢查歸約核心

  

  我們可以看到每個歸約核心執行時間的有序表格。我們還可以看到核心執行了多少次。這在以下幾個方面很有幫助

  • 如果一個核心只佔用極少的時間，例如 0.1%，那麼改進它最多隻能帶來 0.1% 的總體提升。不值得為此投入大量精力。

  • 如果一個核心佔用 2% 的時間，將其效能提高一倍將帶來 1% 的總體提升，這值得投入精力。

對單個 Triton 核心進行基準測試

假設我們想仔細檢視 triton_red_fused\__native_batch_norm_legit_functional_16，這是最耗時的歸約核心，佔前向圖總 wall time 的 2.19%。

我們可以在 fwd.py 中查詢核心名稱，並找到類似如下的註釋

# kernel path: /tmp/torchinductor_shunting/jk/cjk2vm3446xrk7rth7hr6pun7xxo3dnzubwcn6ydrpifal4eykrz.py

為方便起見，我將其重新命名為 k.py。這是一個檔案的貼上內容。

k.py 是一個獨立的 Python 模組，包含核心程式碼及其基準測試。

直接執行 k.py 將報告其執行時間和頻寬

我們可以透過執行以下命令檢查 max-autotune 是否對該核心有幫助

**TORCHINDUCTOR_MAX_AUTOTUNE=1 python /tmp/k.py**

我們還可以臨時新增更多歸約啟發式方法並再次執行指令碼，以檢查這對核心有何幫助。