我們使用 OpenAI Triton 產生部分 GPU 核心。Triton 允許針對特定融合項目產生快速的 GPU 核心,但我們必須針對這類融合調整一些參數。
如果有許多融合,這可能需要很長的時間,因此我們提供一種方法,可在正常執行其他編譯步驟的同時載入這些自動調整結果。自動調整快取仍可在我們進行一些變更後發揮作用:快取中存在的融合會使用快取,而其他融合則會正常自動調整。
建議:快取目錄
--xla_gpu_per_fusion_autotune_cache_dir=your/directory
在指定目錄中使用及維護個別 Fusion Autotune 快取。每個獨立的融合會產生一個檔案。
這種做法的主要優點是,您可以使用相同的快取目錄來執行多次 (採用不同模型) 的 XLA 執行作業,且每次遇到新融合時,快取就會成長,進而加快後續執行的速度。另外,也提供基本支援功能,可同時使用相同快取目錄執行多個 XLA 例項。
XLA 會在需要時讀取現有結果,並在確定後寫入新的結果。
- 在執行 XLA 之前,目錄必須已存在,且必須可寫入。
- 使用者必須自行處理快取無效化:
- 如要從空快取開始,請使用空目錄。
- XLA 版本檢查必須由使用者執行:
- 如果您想針對不同的 XLA 版本使用不同的快取,請使用不同的目錄。
系統會在您未提供參數時,預設關閉快取。
限制:我們不保證可與下文所述的其他快取方法搭配使用。
替代方法:將特定 HLO 的所有結果載入或傾印至單一檔案
您可以使用下列參數傾印/載入自動調整結果:
--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=
如果我們指定 .txt 或 .textproto 檔案,快取會以 textproto 格式傾印,否則會以二進位 protobuf 格式傾印。
在測試中
您也可以在測試中使用已儲存的自動調整功能。如果測試規模很大,特別是測試環境的效能受限時,建議您使用此選項。
只有在自動調整快取包含在測試執行時所使用的同類型 GPU 上產生的結果時,這項功能才能正常運作。
讓測試使用持續性自動調整
我們先假設問題中的測試一律使用相同的 GPU 類型。
我們必須從測試中匯出自動調整結果,例如,將這些參數指定給測試指令:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabled
您必須停用分割功能,才能正確為所有測試取得單一自動調整快取。
接著,我們必須將快取上傳至程式碼存放區。
接著,我們必須將快取新增至測試目標的資料依附元件,並使用環境變數加以載入。
data = ["test_autotune_cache.textproto"], env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" + "$(execpath test_autotune_cache.textproto)"},
(您可以在載入自動調整結果的測試中使用資料分割)。
另請參閱 xla/service/gpu/tests/BUILD 中的測試範例:
- load_autotune_results_using_execpath_test
- load_autotune_results_from_test_workspace_test
- dump_autotune_results_to_test_outputs_test
快取過時
如果對模型進行大量變更,快取有可能不再包含所有融合作業,因此測試的速度會變慢。在這種情況下,我們必須重新產生自動調整快取。
如果我們開始使用新類型的 GPU 來執行測試,同樣適用這項規定。
如果 XLA 編譯器不斷演變並產生不同的融合,快取可能也會過時。