torch.compile 中的編譯時快取¶
建立於：2024 年 6 月 20 日 | 最後更新：2024 年 12 月 06 日 | 最後驗證：2024 年 11 月 05 日
作者： Oguz Ulgen 和 Sam Larsen
簡介¶
PyTorch Inductor 實作了多個快取以減少編譯延遲。此配方示範了如何在 torch.compile 中配置快取的各個部分。
Inductor 快取設定¶
這些快取大部分都在記憶體中，僅在同一進程中使用，並且對使用者是透明的。例外的是儲存編譯後的 FX 圖的快取（FXGraphCache、AOTAutogradCache）。這些快取允許 Inductor 在遇到具有相同張量輸入形狀（和相同配置）的相同圖時，避免跨進程邊界的重新編譯。預設實作將編譯後的成品儲存在系統暫存目錄中。一個可選功能還支援通過將這些成品儲存在 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，此設定啟用遠端的 AOT AutogradCache 功能。目前的實作使用 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¶
此設定啟用 Inductor 自動調諧器的遠端快取。如同遠端 FX graph 快取，目前的實作使用 Redis。1 啟用快取，其他任何值則停用。上述相同的主機/埠環境變數適用於此快取。
TORCHINDUCTOR_FORCE_DISABLE_CACHES¶
將此值設定為 1 以停用所有 Inductor 快取。此設定對於實驗冷啟動編譯時間或強制重新編譯以進行偵錯等任務非常有用。
結論¶
在本指南中，我們了解到 PyTorch Inductor 的快取機制透過利用本地和遠端快取，顯著降低了編譯延遲，這些快取在後台無縫運作，無需使用者干預。此外，我們還探討了各種設定和環境變數，這些設定和變數允許使用者根據其特定需求配置和最佳化這些快取功能。
