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 会导致图中断，并在急切模式（无编译）下运行操作。nn 这适用于操作，而 torch.compiler.disable 适合装饰函数。	如果您使用的是自定义操作，例如 torch.ops.fbgemm.*，导致 torch.compile 函数出现问题，则此 API 非常适合调试和解除阻塞。
torch.compile.allow_in_graph	注释的调用保持在 TorchDynamo 图中。例如，TorchDynamo Dynamo 的黑盒。nn 请注意，AOT 自动微分将对其进行追踪，因此 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 急切模式而不是 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) 将不起作用。