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会中断图形，并以即时执行（不编译）模式运行该操作。 n这对于特定操作是合适的，而`torch.compiler.disable` 适用于装饰函数。	此 API 非常适合用于调试，并且可以帮助解决由自定义操作（如 `torch.ops.fbgemm.*`）引起的 `torch.compile` 函数的问题。
`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 时，TorchDynamo 将拦截并处理每个这样的函数帧（用绿色表示）：

假设函数a_fn与torch.compile存在兼容性问题，并且这部分代码是非关键性的。你可以使用compiler.disable来禁用对函数a_fn的编译。如上所示，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 在相关函数帧包含一些难以支持的功能（如钩子和 autograd.Function）时非常有用，并且你可以确信下游的 PyTorch 组件（例如 AOTAutograd）可以通过装饰的函数安全地进行追踪。当一个函数被装饰为 allow_in_graph 时，TorchDynamo 将其视为黑盒并在生成的图中直接使用。

警告

allow_in_graph 完全跳过装饰的函数中的 TorchDynamo，包括图中断、处理闭包等安全检查。使用allow_in_graph需谨慎，因为 PyTorch 下游组件（如 AOTAutograd）依赖于 TorchDynamo 来处理复杂的 Python 特性，而allow_in_graph会绕过这些特性。这可能导致可靠性和难以调试的问题。

限制

所有的现有 API 都在 TorchDynamo 层级上进行应用。因此，这些 API 只能访问到 TorchDynamo 能看到的信息。这可能导致一些混淆的情况。

例如，torch._dynamo.disallow_in_graph 对 ATen 操作符不起作用，因为它们对 AOT 自动微分是可见的。因此，在上述示例中，torch._dynamo.disallow_in_graph(torch.ops.aten.add) 不会起作用。