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.is_exporting	表示是否通过导出进行图追踪。比 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 特性（如钩子和 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) 将不起作用。