用于细粒度追踪的 TorchDynamo API

注意

在本文档中，torch.compiler.compile 和 torch.compile 可以互换使用。两者在您的代码中都可以正常工作。

torch.compile 对整个用户模型执行 TorchDynamo 跟踪。 然而，可能会有小部分模型代码无法被 torch.compiler 处理。在这种情况下，您可能希望在模型的其余部分运行编译时，禁用该特定部分的编译器。本节描述了用于定义您希望跳过编译的代码部分的现有 API 以及相关的用例。

您可以使用以下表格中列出的API来定义可以在其上禁用编译的代码部分：

TorchDynamo APIs to control fine-grained tracing

API	描述	何时使用？
torch.compiler.disable	禁用装饰函数上的Dynamo以及递归调用的函数。	如果模型的一小部分无法通过torch.compile处理，非常适合用于解封用户。
torch._dynamo.disallow_in_graph	不允许在 TorchDynamo 图中标记的操作。TorchDynamo 会导致图中断，并在急切（不编译）模式下运行操作。nn这对于操作是合适的，而 torch.compiler.disable 适用于装饰函数。	这个API对于调试和解除阻塞都非常出色，如果像torch.ops.fbgemm.*这样的自定义操作导致了torch.compile函数的问题。
torch.compile.allow_in_graph	注释的可调用对象会直接进入 TorchDynamo 图。例如，一个用于 TorchDynamo 的黑色盒子。请注意，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

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) 在上面的例子中将不起作用。