TorchDynamo API 用於細粒度跟蹤

建立日期：2023 年 7 月 28 日 | 最後更新日期：2025 年 6 月 16 日

注意

本文件中，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 會導致圖中斷，並在急切（不編譯）模式下執行運算子。這適用於運算子，而 torch.compiler.disable 適用於裝飾函式。	此 API 非常適合用於除錯和解決自定義運算子（如 torch.ops.fbgemm.*）導致 torch.compile 函數出現問題的情況。
torch.compile.allow_in_graph	註釋的可呼叫物件將按原樣進入 TorchDynamo 圖。例如，對於 TorchDynamo 來說，它是一個黑盒。請注意，AOT Autograd 會對其進行跟蹤，因此 allow_in_graph 僅是 Dynamo 級別的概念。	此 API 對於模型中具有已知 TorchDynamo 難以支援的特性（如 hook 或 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 急切模式而不是 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 能夠提取圖，但隨後您發現下游編譯器出現故障。例如，缺少元核心，或者某個運算子的 Autograd 分派鍵設定不正確。然後，您可以將該運算子標記為 disallow_in_graph，TorchDynamo 將導致圖中斷，並使用 PyTorch 急切模式執行該運算子。

需要注意的是，您必須找到相應的 Dynamo 級運算子，而不是 ATen 級運算子。有關更多資訊，請參閱文件的侷限性部分。

警告

torch._dynamo.disallow_in_graph 是一個全域性標誌。如果您正在比較不同的後端編譯器，則在切換到其他編譯器時，可能需要為被禁止的運算子呼叫 allow_in_graph。

torch.compiler.allow_in_graph

torch.compiler.allow_in_graph 在相關函式幀具有一些已知難以支援的 TorchDynamo 特性（如 hook 和 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) 在上面的示例中將不起作用。