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() 來就地編譯模組，而無需更改其結構。

具體來說，對於在編譯區域內執行的每個 frame，我們將嘗試編譯它，並將編譯結果緩存在程式碼物件上以供將來使用。如果先前的編譯結果不適用於後續呼叫（這稱為“guard failure”），則單個 frame 可能會被編譯多次，您可以使用 TORCH_LOGS=guards 來除錯這些情況。最多可以將 torch._dynamo.config.cache_size_limit 個編譯結果與一個 frame 關聯，預設值為 8；此時我們將退回到 eager 模式。請注意，編譯快取是每個程式碼物件，而不是 frame；如果您動態建立函數的多個副本，它們將共享同一個程式碼快取。

參數

• model (Callable) – 要優化的模組/函數

• fullgraph (bool) – 如果為 False（預設值），torch.compile 會嘗試發現函數中可編譯的區域，以便進行優化。如果為 True，則我們要求將整個函數捕獲到單個圖中。如果這不可能（也就是說，如果存在圖中斷），則會引發錯誤。

• dynamic (bool 或 None) – 使用動態形狀追蹤。當此值為 True 時，我們將預先嘗試產生一個盡可能動態的 kernel，以避免在大小更改時重新編譯。這可能並不總是有效，因為某些操作/優化將強制進行專業化；使用 TORCH_LOGS=dynamic 來除錯過度專業化。當此值為 False 時，我們將永遠不會產生動態 kernel，我們將始終進行專業化。預設情況下（None），我們會自動檢測是否發生了動態變化，並在重新編譯時編譯更動態的 kernel。

• backend (str 或 Callable) –

  要使用的後端。

  • ”inductor” 是預設的後端，在效能和額外負擔之間取得良好的平衡。

  • 可以使用 torch._dynamo.list_backends() 查看非實驗性的內部後端。

  • 可以使用 torch._dynamo.list_backends(None) 查看實驗性或除錯用的內部後端。

  • 要註冊外部自定義後端：https://pytorch.dev.org.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 graphs 降低 Python 額外負擔的模式，適用於小批量。降低額外負擔可能會增加記憶體用量，因為我們會快取調用所需的工作區記憶體，以便在後續執行中不必重新分配。不保證降低額外負擔一定有效；目前，我們只降低 CUDA only graphs 的額外負擔，而且不會修改輸入。在其他情況下，CUDA graphs 可能不適用；請使用 TORCH_LOG=perf_hints 進行除錯。

  • ”max-autotune” 是一種在支援的裝置上利用 Triton 或基於樣板的矩陣乘法，以及在 GPU 上利用基於 Triton 的卷積的模式。預設情況下，它會在 GPU 上啟用 CUDA graphs。

  • ”max-autotune-no-cudagraphs” 是一種與 “max-autotune” 類似，但不使用 CUDA graphs 的模式。

  • 要查看每種模式設定的確切配置，您可以呼叫 torch._inductor.list_mode_options()。

• options (dict) –

  要傳遞給後端的選項字典。一些值得嘗試的選項是：

  • epilogue_fusion，將逐點運算融合到樣板中。需要同時設定 max_autotune。

  • max_autotune，將進行分析以選擇最佳的矩陣乘法配置。

  • fallback_random，在除錯準確性問題時很有用。

  • shape_padding，填充矩陣形狀以更好地對齊 GPU 上的載入，尤其適用於 tensor cores。

  • triton.cudagraphs，將使用 CUDA graphs 降低 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)