TorchDynamo API 用于细粒度跟踪

创建于：2023 年 7 月 28 日 | 最后更新于：2025 年 6 月 16 日

注意

本文档中 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 会导致图中断，并以惰性（无编译）模式运行该操作。nn这适用于操作，而 torch.compiler.disable 适用于装饰函数。	此 API 对于调试和解决自定义操作（如 torch.ops.fbgemm.*）导致 torch.compile 函数出现问题非常有用。
torch.compile.allow_in_graph	被注解的可调用对象按原样放入 TorchDynamo 图中。例如，对于 TorchDynamo 的黑盒。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) 将不起作用。