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（無編譯）模式下執行運算。nn 這適用於運算，而 torch.compiler.disable 適用於裝飾函式。	如果自訂運算（例如 torch.ops.fbgemm.*）導致 torch.compile 函式出現問題，則此 API 非常適合用於偵錯和解除封鎖。
torch.compile.allow_in_graph	已標註的可呼叫物件會照原樣出現在 TorchDynamo 圖形中。例如，TorchDynamo Dynamo 的黑盒子。nn 請注意，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.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 能夠提取圖形，但隨後您會看到下游編譯器失敗。例如，缺少中繼內核，或者某些 Autograd 調度鍵針對特定運算子設置不正確。然後，您可以將該運算子標記為 disallow_in_graph，TorchDynamo 將導致圖形中斷，並使用 PyTorch 急切模式運行該運算子。

需要注意的是，您必須找到對應的 Dynamo 級別運算子，而不是 ATen 級別運算子。請參閱文檔的限制部分了解更多信息。

警告

torch._dynamo.disallow_in_graph 是一個全局標誌。如果您要比較不同的後端編譯器，則在切換到其他編譯器時，可能需要為不允許的運算子調用 allow_in_graph。

torch.compiler.allow_in_graph

當相關的函數框架具有一些已知難以支持的 TorchDynamo 功能（例如鉤子和 autograd.Function）時，torch.compiler.allow_in_graph 很有用，並且您確信下游 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) 在上述範例中將無法使用。