編譯時快取配置¶
作者: Oguz Ulgen 和 Sam Larsen
引言¶
PyTorch 編譯器實現了多種快取以減少編譯延遲。本示例集展示瞭如何在 torch.compile 中配置快取的不同部分。
Inductor 快取設定¶
大多數這些快取是記憶體中的,僅在同一程序內使用,並且對使用者是透明的。一個例外是儲存編譯後的 FX 圖的快取(FXGraphCache, AOTAutogradCache)。當 Inductor 遇到具有相同 Tensor 輸入形狀(和相同配置)的同一圖時,這些快取可以避免跨程序邊界的重新編譯。預設實現將編譯後的工件儲存在系統臨時目錄中。一個可選功能還支援透過將這些工件儲存在 Redis 資料庫中來在叢集內共享它們。
有一些設定與快取,特別是 FX 圖快取相關。這些設定可以透過下面列出的環境變數訪問,或硬編碼到 Inductor 的配置檔案中。
TORCHINDUCTOR_FX_GRAPH_CACHE¶
此設定啟用本地 FX 圖快取功能,該功能將工件儲存在主機的臨時目錄中。將其設定為 1 會啟用該功能,任何其他值則會停用它。預設情況下,磁碟位置是按使用者名稱區分的,但使用者可以透過指定 TORCHINDUCTOR_CACHE_DIR(見下文)來啟用跨使用者名稱共享。
TORCHINDUCTOR_AUTOGRAD_CACHE¶
此設定擴充套件了 FXGraphCache,以便在 AOTAutograd 級別而非 Inductor 級別儲存快取結果。將其設定為 1 會啟用此功能,任何其他值則會停用它。預設情況下,磁碟位置按使用者名稱區分,但使用者可以透過指定 TORCHINDUCTOR_CACHE_DIR(見下文)來啟用跨使用者名稱共享。TORCHINDUCTOR_AUTOGRAD_CACHE 需要 TORCHINDUCTOR_FX_GRAPH_CACHE 才能工作。相同的快取目錄儲存 AOTAutogradCache 的快取條目(位於 {TORCHINDUCTOR_CACHE_DIR}/aotautograd 下)和 FXGraphCache 的快取條目(位於 {TORCHINDUCTOR_CACHE_DIR}/fxgraph 下)。
TORCHINDUCTOR_CACHE_DIR¶
此設定指定所有磁碟上快取的位置。預設情況下,位置在系統臨時目錄下的 torchinductor_<username> 中,例如 /tmp/torchinductor_myusername。
請注意,如果在環境中未設定 TRITON_CACHE_DIR,Inductor 會將 Triton 快取目錄設定為此臨時位置,在其 Triton 子目錄下。
TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE¶
此設定啟用遠端 FX 圖快取功能。當前實現使用 Redis。設定為 1 啟用快取,任何其他值則停用它。以下環境變數配置 Redis 伺服器的主機和埠:
TORCHINDUCTOR_REDIS_HOST(預設為 localhost) TORCHINDUCTOR_REDIS_PORT(預設為 6379)
注意
請注意,如果 Inductor 找到了一個遠端快取條目,它會將編譯後的工件儲存在本地磁碟快取中;後續在同一臺機器上執行時,將直接使用該本地工件。
TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE¶
類似於 TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE,此設定啟用遠端 AOTAutogradCache 功能。當前實現使用 Redis。將其設定為 1 會啟用快取,任何其他值則停用它。以下環境變數用於配置 Redis 伺服器的主機和埠:* TORCHINDUCTOR_REDIS_HOST(預設為 localhost)* TORCHINDUCTOR_REDIS_PORT(預設為 6379)
TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE` 需要啟用 TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE 才能工作。同一個 Redis 伺服器可以用於儲存 AOTAutograd 和 FXGraph 快取結果。
TORCHINDUCTOR_AUTOTUNE_REMOTE_CACHE¶
此設定啟用 TorchInductor 的自動調優器的遠端快取。類似於遠端 FX 圖快取,當前實現使用 Redis。將其設定為 1 會啟用快取,任何其他值則停用它。上面提到的相同的主機/埠環境變數也適用於此快取。
TORCHINDUCTOR_FORCE_DISABLE_CACHES¶
將此值設定為 1 以停用所有 Inductor 快取。此設定對於實驗冷啟動編譯時間或出於除錯目的強制重新編譯等任務非常有用。
結論¶
在本示例集中,我們學習瞭如何配置 PyTorch 編譯器的快取機制。此外,我們還探索了各種設定和環境變數,這些設定和環境變數允許使用者根據其特定需求配置和最佳化這些快取功能。