torch.compile

torch.compile(model: Callable[[_InputT], _RetT], *, fullgraph: bool = False, dynamic: Optional[bool] = None, backend: Union[str, Callable] = 'inductor', mode: Optional[str] = None, options: Optional[dict[str, Union[str, int, bool]]] = None, disable: bool = False) → Callable[[_InputT], _RetT]

torch.compile(model: None = None, *, fullgraph: bool = False, dynamic: Optional[bool] = None, backend: Union[str, Callable] = 'inductor', mode: Optional[str] = None, options: Optional[dict[str, Union[str, int, bool]]] = None, disable: bool = False) → Callable[[Callable[[_InputT], _RetT]], Callable[[_InputT], _RetT]]

使用 TorchDynamo 和指定的後端最佳化給定的模型/函式。如果要編譯 torch.nn.Module，還可以使用 torch.nn.Module.compile() 就地（inplace）編譯模組，而無需改變其結構。

具體而言，對於在編譯區域內執行的每個幀（frame），我們將嘗試對其進行編譯，並將編譯結果快取到程式碼物件（code object）上以備將來使用。如果之前的編譯結果不適用於後續呼叫（這被稱為“守衛失敗”或“guard failure”），單個幀可能會被編譯多次。您可以使用 TORCH_LOGS=guards 來除錯這些情況。一個幀最多可以關聯 torch._dynamo.config.recompile_limit 個編譯結果，預設值為 8；達到此限制後，我們將回退到 eager 模式。請注意，編譯快取是基於程式碼物件（code object）而非幀（frame）的；如果您動態建立函式的多個副本，它們將共享相同的程式碼快取。

引數

• model (Callable) – 要最佳化的模組/函式

• fullgraph (bool) – 如果為 False（預設值），torch.compile 會嘗試在函式中發現可編譯的區域進行最佳化。如果為 True，則要求整個函式必須能被捕獲到一個單獨的計算圖中。如果無法做到（即存在圖斷點graph breaks），則會引發錯誤。

• dynamic (bool or None) – 使用動態形狀追蹤。當此引數為 True 時，我們將提前嘗試生成儘可能動態的核心，以避免在形狀變化時重新編譯。但這並非總是有效，因為某些操作/最佳化會強制進行特化；使用 TORCH_LOGS=dynamic 來除錯過度特化問題。當此引數為 False 時，我們將永遠不會生成動態核心，而是始終進行特化。預設值（None）下，我們自動檢測是否發生了動態行為，並在重新編譯時編譯更動態的核心。

• backend (str or Callable) –

  要使用的後端

  • “inductor” 是預設後端，它在效能和開銷之間取得了很好的平衡

  • 非實驗性的內建後端可以使用 torch._dynamo.list_backends() 檢視

  • 實驗性或除錯用的內建後端可以使用 torch._dynamo.list_backends(None) 檢視

  • 註冊外部自定義後端：https://pytorch.com.tw/docs/main/torch.compiler_custom_backends.html#registering-custom-backends

• mode (str) –

  可以是 “default”、“reduce-overhead”、“max-autotune” 或 “max-autotune-no-cudagraphs” 之一

  • “default” 是預設模式，它在效能和開銷之間取得了很好的平衡

  • “reduce-overhead” 模式使用 CUDA 圖減少 Python 的開銷，對小批次資料很有用。減少開銷可能會以增加記憶體使用為代價，因為我們將快取呼叫所需的 workspace 記憶體，以便後續執行時不必重新分配。減少開銷並非始終有效；目前，我們只對不修改輸入的純 CUDA 圖減少開銷。在其他一些情況下，CUDA 圖不適用；可以使用 TORCH_LOG=perf_hints 進行除錯。

  • “max-autotune” 模式在支援的裝置上利用 Triton 或基於模板的矩陣乘法，並在 GPU 上利用基於 Triton 的卷積。它在 GPU 上預設啟用 CUDA 圖。

  • “max-autotune-no-cudagraphs” 模式類似於 “max-autotune”，但不使用 CUDA 圖

  • 要檢視每種模式設定的具體配置，可以呼叫 torch._inductor.list_mode_options()

• options (dict) –

  傳遞給後端的可選項字典。一些值得嘗試的重要選項包括：

  • epilogue_fusion 將逐點（pointwise）操作融合到模板中。需要同時設定 max_autotune

  • max_autotune 將透過 profile 選擇最佳的 matmul 配置

  • fallback_random 在除錯精度問題時很有用

  • shape_padding 填充矩陣形狀，以更好地對齊 GPU 上的載入，特別是對於張量核心（tensor cores）

  • triton.cudagraphs 使用 CUDA 圖減少 Python 的開銷

  • trace.enabled 是最有用的除錯標誌之一

  • trace.graph_diagram 將顯示融合後的計算圖影像

  • 對於 inductor 後端，可以透過呼叫 torch._inductor.list_options() 檢視其支援的完整配置列表

• disable (bool) – 將 torch.compile() 變成空操作（no-op）用於測試

示例

@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
def foo(x):
    return torch.sin(x) + torch.cos(x)