TunableOp#
創建於: 2024年6月3日 | 最後更新於: 2025年6月13日
概述#
此模組公開了 TunableOp 介面。
某些操作，例如 GEMMs，可以使用多個庫或多個技術來實現。例如，GEMM 可以為 CUDA 或 ROCm 實現，使用 blas 或 blasLt 庫。此外，ROCm 的 rocblas 和 hipblaslt 庫允許使用者查詢所有可能的演算法，然後選擇一個。如何知道哪種實現是最快的並應被選擇？這正是 TunableOp 提供的。
單獨啟用 TunableOp 和調優#
TunableOp 功能的啟用與調優階段本身的啟用是分開的。啟用 TunableOp 意味著 PyTorch 將用其 Tunable 實現替換任何標準運算子。任何對 TunableOp 的呼叫首先會檢查是否已針對給定的運算子輸入進行了調優。如果是，它將立即呼叫已調優的操作；即使啟用了調優設定，也不會進行進一步的調優。相反，如果未找到調優結果，並且啟用了調優，TunableOp 將為給定的一組輸入對該運算子的每個已註冊實現進行基準測試，並選擇最快的。
檔案輸入和輸出#
首次呼叫任何 TunableOp 時，內部的已調優操作資料庫將透過嘗試從給定檔案中讀取結果來準備。預設檔名是 'tunableop_results.csv'。為了支援在多個程序中使用多個 GPU 進行調優，GPU 裝置序數會自動插入到檔名中，以避免多個程序覆蓋同一檔案。
如果啟用了調優，並且在工作負載過程中發現了新的調優項，它還將寫入到同一個檔名中，包含所有調優項，包括啟動時讀取的以及執行時發現的新項。例如，這可以用於透過重用同一個檔案來跨多個工作負載構建調優檔案。輸出檔案在應用程式終止時自動建立。此行為可以透過 C++ 和 Python API 控制，但不能透過環境變數控制。
假設您指定了一個檔名，您將得到一個包含類似內容的 CSV 檔案
Validator,PT_VERSION,2.2.0
Validator,ROCM_VERSION,6.0.0.0-12969-1544e39
Validator,HIPBLASLT_VERSION,0.6.0-a9c5cc7
Validator,ROCBLAS_VERSION,4.0.0-72e57364-dirty
GemmTunableOp_float_NT,nt_25088_4096_64,Gemm_Hipblaslt_1219,1.262
GemmTunableOp_float_NT,nt_4096_4096_64,Gemm_Rocblas_1216,0.033
注意“驗證器”行。如果您更改了庫版本、ROCm 版本或 PyTorch 版本，TunableOp 將檢測到這一點並拒絕該調優檔案，因為先前的調優可能受到其他軟體更改的影響。
剩餘的行是您在執行過程中遇到的每個 TunableOp 的已調優解決方案。每行包含 4 個逗號分隔的欄位：運算子名稱、運算子引數、解決方案名稱和平均執行時間。執行時間是一個可選欄位。CSV 檔案可以編輯，但需謹慎。例如，解決方案名稱（欄位 3）可以更改為“Default”，它將回退到原始 PyTorch 未調優的實現。或者，對於 ROCm 的 hipBLAS 或 hipBLASLt 庫，如果您知道特定的解決方案索引，可以透過替換值來覆蓋 TunableOp 選擇的解決方案。運算子名稱和引數（欄位 1 和 2）是內部命名的，不應修改。對於 GemmTunableOp，欄位 1 指示資料型別以及輸入是否已轉置 (T) 或未轉置 (N)，欄位 2 指示 M、N、K 輸入形狀。
有一個選項可以啟用詳細輸出，但僅推薦用於除錯目的。這將產生大量診斷訊息，但可能有助於檢視 TunableOp 是否被使用。否則，TunableOp 完全靜默，除了檔案輸出，除非在使用過程中出現警告或錯誤。詳細選項僅透過設定環境變數 PYTORCH_TUNABLEOP_VEROBSE=1 來可用。
關於調優行為、預熱和快取效應的說明#
調優一個運算子包括遍歷已註冊的實現列表並對每個實現進行效能分析。效能分析是透過將單個實現迴圈執行多次並取平均執行時間來建立的。在調優之前還有一個可選的預熱階段，可以幫助硬體達到穩定的功率狀態。在工作負載調優期間，各種硬體快取更有可能產生命中，而不是在不調優時。有一些選項可以重新整理指令快取和旋轉輸入張量，這可能有助於更準確地分析已調優的運算子，就像該運算子在更大的工作負載中執行一樣，而不是在一個緊湊、重複的迴圈中。
預設情況下，對於給定的運算子，每個可能的解決方案將執行 100 次迭代，或者執行 30ms 內可以執行的儘可能多的迭代，取兩者中的較小值，並計算其平均執行時間。在所有成功進行效能分析的解決方案中，將選擇最快的解決方案。如果給定的解決方案未達到與預設實現相同的精度或解決方案返回了錯誤程式碼，效能分析可能會失敗。
當前可調優的運算子#
ROCm 的 TunableGemm#
目前只實現了一個 ROCm 的 TunableGemm。請注意，PyTorch 的 CUDA 構建在使用 TunableOp 時可以正常工作，但對 CUDA 構建可用的唯一解決方案是“Default”實現，即透過 TunableOp 呼叫的原始 cuBLAS 預設實現。任何對 at::cuda::blas::gemm() 或 ::bgemm() 的呼叫在啟用時都將透過 TunableOp 進行路由。呼叫具有給定輸入引數集（transa、transb、m、n、k）的 gemm() 將嘗試使用 rocblas 和 hipblaslt 中可用的最快實現。
離線調優#
動機#
離線調優有幾種用例。
一種用例涉及高記憶體利用率的工作負載，此時常規調優可能導致記憶體不足。
另一種用例是計算密集型工作負載。在這種情況下，一次收集工作負載的 GEMMs，然後用不同的調優引數或庫重複調優，會更節省資源。
工作流程#
基本上有兩個步驟：1) 設定環境變數以收集未調優的 GEMM，這將生成 tunableop_untuned0.csv
export PYTORCH_TUNABLEOP_ENABLED=1
export PYTORCH_TUNABLEOP_TUNING=0
export PYTORCH_TUNABLEOP_RECORD_UNTUNED=1
...
執行一個 Python 指令碼，該指令碼讀取
tunableop_untuned0.csv並生成tunableop_results0.csv，如下所示
import torch.cuda.tunable as tunable
import os
os.putenv("PYTORCH_TUNABLEOP_ENABLED", "1")
os.putenv("PYTORCH_TUNABLEOP_TUNING", "1")
os.putenv("PYTORCH_TUNABLEOP_RECORD_UNTUNED", "0")
tunable.tune_gemm_in_file("tunableop_untuned0.csv")
也可以同時處理多個未調優檔案，並將 GEMMs 分佈到單個節點中的多個 GPU 上進行調優。第一步，首先收集 GEMMs 並消除重複的 GEMMs。接下來，將 GEMMs 分發到不同的 GPU 上進行調優。在所有 GEMMs 都調優完成後，來自所有 GPU 的結果將被收集到一個檔案中，該檔案的基本檔名將附加 _full0（例如 tunableop_results_full0.csv）。最後，這個包含收集結果的新檔案將被複制 N 次，每次複製一個給一個 GPU，以方便使用者在 N 個 GPU 上執行具有已調優配置的工作負載。
if __name__ == "__main__":
num_gpus = 8 # number of GPUs that will be used during the tuning process
tunable.mgpu_tune_gemm_in_file("tunableop_untuned?.csv", num_gpus)
請注意，mgpu_tune_gemm_in_file API 的用法與其單 GPU 對應項（tune_gemm_in_file）不同。呼叫該 API 的 Python 指令碼的主體必須如所示用 main() 包裹起來，因為使用了 concurrent.futures 模組。傳遞給 mgpu_tune_gemm_in_file 的引數必須包含一個萬用字元表示式（? 或 *）來生成要處理的 GEMMs 的未調優檔案列表。num_gpus 必須介於 1 和可用 GPU 總數之間。
API 參考#
- torch.cuda.tunable.tuning_enable(val=True)[source]#
啟用 TunableOp 實現的調優。
啟用後，如果找不到已調優的條目，將執行調優步驟並記錄該條目。
- torch.cuda.tunable.record_untuned_enable(val=True)[source]#
為離線調優啟用 TunableOp 操作的未調優記錄。
啟用後，如果找不到已調優的條目，則將其寫入未調優檔案。
- torch.cuda.tunable.set_max_tuning_duration(duration)[source]#
設定調優給定解決方案的最大時間（毫秒）。
如果同時設定了最大調優時間和迭代次數，將遵守兩者中的較小值。至少會始終執行 1 次調優迭代。
- torch.cuda.tunable.set_max_tuning_iterations(iterations)[source]#
設定調優給定解決方案的最大迭代次數。
如果同時設定了最大調優時間和迭代次數，將遵守兩者中的較小值。至少會始終執行 1 次調優迭代。
- torch.cuda.tunable.set_filename(filename, insert_device_ordinal=False)[source]#
設定用於輸入/輸出調優結果的檔名。
如果
insert_device_ordinal為True，則當前裝置序數將自動新增到給定的檔名中。這可以在每個 GPU 一個程序的場景中使用，以確保所有程序寫入單獨的檔案。
- torch.cuda.tunable.write_file_on_exit(val)[source]#
在上下文銷燬時，將檔案寫入磁碟。
這有助於在應用程式因正常操作或錯誤而終止時將結果最終重新整理到磁碟。可以透過手動呼叫
write_file()來手動重新整理結果。
- torch.cuda.tunable.write_file(filename=None)[source]#
將結果寫入 CSV 檔案。
如果未給出
filename，則呼叫get_filename()。- 返回型別
- torch.cuda.tunable.read_file(filename=None)[source]#
從 TunableOp CSV 檔案讀取結果。
如果未給出
filename，則呼叫get_filename()。- 返回型別
- torch.cuda.tunable.mgpu_tune_gemm_in_file(filename_pattern, num_gpus)[source]#
處理一個或多個檔案並將工作分配到一或多個 GPU 上。
