TorchDynamo 用於細粒度跟蹤的 API

注意

在本文件中，torch.compiler.compile 和 torch.compile 可互換使用。這兩個版本在你的程式碼中都可以正常工作。

torch.compile 對整個使用者模型執行 TorchDynamo 跟蹤。但是，模型程式碼的一小部分可能無法由 torch.compiler 處理。在這種情況下，你可能希望在該特定部分停用編譯器，而在模型的其餘部分執行編譯。本節描述了用於定義你希望跳過編譯的程式碼部分以及相關用例的現有 API。

下表列出了可用於定義要停用編譯的程式碼部分的 API

用於控制細粒度跟蹤的 TorchDynamo API

API	描述	何時使用？
torch.compiler.disable	停用對裝飾函式及其遞迴呼叫函式的 Dynamo 編譯。	如果模型的一小部分無法由 torch.compile 處理，這對解決使用者問題非常有用。
torch._dynamo.disallow_in_graph	禁止在 TorchDynamo 圖中出現標記的運算元。TorchDynamo 會導致圖中斷，並在 eager（無編譯）模式下執行該運算元。 這適用於運算元，而 torch.compiler.disable 適用於裝飾函式。	如果像 torch.ops.fbgemm.* 這樣的自定義運算元導致 torch.compile 函數出現問題，此 API 對於除錯和解決問題都非常有用。
torch.compile.allow_in_graph	帶有註解的可呼叫物件會原樣進入 TorchDynamo 圖中。例如，對於 TorchDynamo Dynamo 來說，它是一個黑盒。 請注意，AOT Autograd 會跟蹤它，因此 allow_in_graph 僅是 Dynamo 級別的概念。	此 API 對於模型中包含已知 TorchDynamo 難以支援的特性（例如鉤子或 autograd.Function）的部分非常有用。但是，每次使用 allow_in_graph 都必須經過仔細篩選（無圖中斷、無閉包）。
torch._dynamo.graph_break	新增一個圖中斷。圖中斷前後的程式碼會經過 TorchDynamo 處理。	在部署中很少有用 - 如果你認為需要此功能，很可能你需要 disable 或 disallow_in_graph。
torch.compiler.is_compiling	指示圖是否作為 torch.compile() 或 torch.export() 的一部分被執行/跟蹤。
torch.compiler.is_dynamo_compiling	指示圖是否透過 TorchDynamo 跟蹤。它比 torch.compiler.is_compiling() 標誌更嚴格，只有在使用 TorchDynamo 時才會被設定為 True。
torch.compiler.is_exporting	指示圖是否透過 export 跟蹤。它比 torch.compiler.is_compiling() 標誌更嚴格，只有在使用 torch.export 時才會被設定為 True。

torch.compiler.disable

torch.compiler.disable 停用對裝飾函式幀以及從裝飾函式幀遞迴呼叫的所有函式幀的編譯。

TorchDynamo 會攔截每個 Python 函式幀的執行。所以，假設你有一個程式碼結構（如下圖所示），其中函式 fn 呼叫函式 a_fn 和 b_fn。而 a_fn 呼叫 aa_fn 和 ab_fn。當你使用 PyTorch eager 模式而不是 torch.compile 時，這些函式幀會按原樣執行。使用 torch.compile 時，TorchDynamo 會攔截這些函式幀中的每一個（由綠色表示）

假設函式 a_fn 導致 torch.compile 出現問題。並且這部分模型並不關鍵。你可以在函式 a_fn 上使用 compiler.disable。如上所示，TorchDynamo 將停止檢視源自 a_fn 呼叫的幀（白色表示原始 Python 行為）。

要跳過編譯，你可以用 @torch.compiler.disable 裝飾有問題的函式。

如果你不想更改原始碼，也可以使用非裝飾器語法。但是，如果可能，我們建議你避免使用這種風格。在這種情況下，你必須確保原始函式的所有使用者現在都在使用打了補丁的版本。

torch._dynamo.disallow_in_graph

torch._dynamo.disallow_in_graph 禁止某個運算元（而非函式）出現在 TorchDynamo 提取的圖中。請注意，這適用於運算元，而不適用於像 _dynamo.disable 那樣的通用函式。

假設你使用 PyTorch 編譯模型。TorchDynamo 能夠提取一個圖，但下游編譯器隨後失敗了。例如，缺少 meta kernel，或者某個運算元的 Autograd dispatch key 設定不正確。這時你可以將該運算元標記為 disallow_in_graph，TorchDynamo 將導致圖中斷，並使用 PyTorch eager 模式執行該運算元。

需要注意的是，你必須找到對應的 Dynamo 級別運算元，而不是 ATen 級別運算元。更多資訊請參閱文件的限制部分。

警告

torch._dynamo.disallow_in_graph 是一個全域性標誌。如果你正在比較不同的後端編譯器，切換到另一個編譯器時，可能需要為被停用的運算元呼叫 allow_in_graph。

torch.compiler.allow_in_graph

torch.compiler.allow_in_graph 在相關函式幀包含已知 TorchDynamo 難以支援的特性（例如鉤子和 autograd.Function）時非常有用，並且你確信下游 PyTorch 元件（如 AOTAutograd）可以安全地跟蹤裝飾函式。當函式被 allow_in_graph 裝飾時，TorchDynamo 將其視為黑盒，並將其原樣放入生成的圖中。

警告

allow_in_graph 完全跳過對裝飾函式的 TorchDynamo 處理，忽略所有 TorchDynamo 安全檢查，包括圖中斷、閉包處理等。請謹慎使用 allow_in_graph。PyTorch 的下游元件（如 AOTAutograd）依賴 TorchDynamo 來處理複雜的 Python 特性，但 allow_in_graph 繞過了 TorchDynamo。使用 allow_in_graph 可能會導致健全性問題和難以除錯的問題。

限制

所有現有 API 都應用於 TorchDynamo 級別。因此，這些 API 只能看到 TorchDynamo 看到的內容。這可能導致令人困惑的情況。

例如，torch._dynamo.disallow_in_graph 對 ATen 運算元無效，因為它們對 AOT Autograd 可見。例如，在上面的例子中，torch._dynamo.disallow_in_graph(torch.ops.aten.add) 將無效。