TorchInductor GPU 分析

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

相關環境變數

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

• TORCHINDUCTOR_UNIQUE_KERNEL_NAMES

  • 預設情況下，TorchInductor 將 Triton 核心命名為 ‘triton\_’。啟用此環境變數後，inductor 會在追蹤中產生更有意義的核心名稱，例如 triton_poi_fused_cat_155，其中包含核心類別（pointwise 的 poi）和原始的 ATen 運算子。預設情況下，此配置已停用，以提高編譯快取命中的機率。

• TORCHINDUCTOR_BENCHMARK_KERNEL

  • 啟用此選項將使 inductor codegen harness 能夠對個別的 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 行這樣的行，一個用於 forward 圖，另一個用於 backward 圖。

對於我們的範例命令，我們分別為 forward 和 backward 圖取得以下編譯後的模組

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

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

3. 現在我們可以深入研究每個個別編譯模組的效能。讓我們選擇 forward 圖進行說明。為了方便起見，我將其命名為 fwd.py。使用 -p 參數直接執行它

   **> python fwd.py -p**
   

請參閱此範例 gist 中的完整輸出日誌。

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

• 我們為該 profile 寫入一個 chrome 追蹤檔案，因此我們可以載入追蹤並與之互動。在日誌中，尋找如下所示的行以找到追蹤檔案的路徑。

Chrome 追蹤 profile 已寫入 /tmp/compiled_module_profile.json

將追蹤載入到 Chrome 中（在 chrome 瀏覽器中造訪 chrome://tracing 並按照 UI 建議載入檔案）將顯示如下 UI

您可以放大和縮小以查看 profile。

• 我們透過類似於以下日誌行的訊息來報告相對於 wall time 的 GPU 時間百分比

  Percent of time when GPU is busy: 102.88% (GPU 忙碌的時間百分比：102.88%)

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

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

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

  這表示模型有大量的 CPU 額外開銷。這與啟用 cudagraphs 可以大大提高 densenet121 的效能這一事實一致。

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

  • pointwise 核心佔用 28.58%

  • reduction 核心佔用 13.85%

  • persistent reduction 核心佔用 3.89%

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

  此資訊可在每個核心類別的報告摘要行（最後一行）中找到。

• 我們也可以放大到某個類別的核心。例如，讓我們檢查 reduction 核心

  

  我們可以看到每個個別 reduction 核心的執行時間的排序表格。我們還可以看到核心的執行次數。這在幾個方面很有幫助

  • 如果核心僅佔用少量時間，例如 0.1%，則改進它最多只能帶來 0.1% 的總體增益。不值得在這方面投入大量精力。

  • 如果核心佔用 2% 的時間，將其提高 2 倍將帶來 1% 的總體增益，這證明了這種努力是值得的。

基準測試個別 Triton 核心

假設我們要更仔細地查看 triton_red_fused\__native_batch_norm_legit_functional_16，它是最昂貴的 reduction 核心，並且佔 forward 圖總體 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**

我們也可以暫時添加更多的 reduction 啟發法，然後再次執行該腳本以檢查這對核心有何幫助。