編譯時快取配置#
作者： Oguz Ulgen 和 Sam Larsen
簡介#
PyTorch Compiler 實現了多個快取以減少編譯延遲。本示例演示瞭如何在 torch.compile 中配置快取的各個部分。
先決條件#
在開始此秘籍之前，請確保您已具備以下條件
對
torch.compile的基本瞭解。參見PyTorch 2.4 或更高版本
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 Compiler 的快取機制。此外，我們還探討了允許使用者根據其特定需求配置和最佳化這些快取功能的各種設定和環境變數。
