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 難以支援的功能的部分，例如 hooks 或 autograd.Function。但是，每次使用 allow_in_graph **都必須仔細檢查** (沒有圖形中斷，沒有 closures)。
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.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

當相關的函式框架具有一些已知的難以支援的 TorchDynamo 功能 (例如 hooks 和 autograd.Function)，並且您確信下游 PyTorch 元件 (例如 AOTAutograd) 可以安全地追蹤被裝飾的函式時，torch.compiler.allow_in_graph 很有用。 當使用 allow_in_graph 裝飾函式時，TorchDynamo 會將其視為黑盒子，並按原樣將其放入產生的圖形中。

警告

allow_in_graph 會完全跳過被裝飾的函式上的 TorchDynamo，並省略所有 TorchDynamo 安全檢查，包括圖形中斷、處理 closures 等。 請謹慎使用 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) 在上面的範例中將不起作用。