持續自動調整 (僅限 GPU)

我們使用 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 類型。

  1. 我們必須從測試中匯出自動調整結果,例如,將這些參數指定給測試指令:

    --test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto
    --test_sharding_strategy=disabled
    

    您必須停用分割功能,才能正確為所有測試取得單一自動調整快取。

  2. 接著,我們必須將快取上傳至程式碼存放區。

  3. 接著,我們必須將快取新增至測試目標的資料依附元件,並使用環境變數加以載入。

    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 編譯器不斷演變並產生不同的融合,快取可能也會過時。