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（无编译）模式下运行该算子。 这适用于算子，而 torch.compiler.disable 适用于装饰函数。	如果像 torch.ops.fbgemm.* 这样的自定义算子导致 torch.compile 函数出现问题，此 API 对于调试和解决问题都非常有用。
torch.compile.allow_in_graph	带有注解的可调用对象会原样进入 TorchDynamo 图中。例如，对于 TorchDynamo Dynamo 来说，它是一个黑盒。 请注意，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	指示图是否通过 export 跟踪。它比 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 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 能够提取一个图，但下游编译器随后失败了。例如，缺少 meta kernel，或者某个算子的 Autograd dispatch key 设置不正确。这时你可以将该算子标记为 disallow_in_graph，TorchDynamo 将导致图中断，并使用 PyTorch eager 模式运行该算子。

需要注意的是，你必须找到对应的 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) 将无效。