常见问题

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

作者: Mark Saroufim

torch.compile 支持训练吗？

torch.compile 支持训练，使用 AOTAutograd 来捕获反向传播。

1. .forward() 图和 `optimizer.step()` 由 TorchDynamo 的 Python `evalframe` 前端捕获。

2. 对于 TorchDynamo 捕获的 `forward()` 的每个片段，它都使用 AOTAutograd 生成一个反向传播图片段。

3. 前向和后向图的每对（可选地）通过最小割进行分区，以在前向和后向之间保存最小状态。

4. 前向和后向对被封装在 `autograd.function` 模块中。

5. 用户代码调用 `.backward()` 仍然会触发 eager 的 autograd 引擎，该引擎将每个 *编译过的反向传播* 图当作一个操作来运行，同时也会运行任何未编译的 eager 操作的 `.backward()` 函数。

您支持分布式代码吗？

torch.compile 支持 `DistributedDataParallel` (DDP)。其他分布式训练库的支持正在考虑中。

分布式代码对于 Dynamo 来说具有挑战性的主要原因是，AOTAutograd 会展开前向和后向传播，并为后端提供 2 个图进行优化。这对于分布式代码来说是一个问题，因为我们希望理想情况下将通信操作与计算重叠。Eager PyTorch 通过各种方式实现这一点，例如使用 autograd hook、module hook 以及修改/突变 module 状态。在 Dynamo 的粗暴应用中，本应在反向传播期间在操作之后立即运行的 hook 可能会被延迟到整个反向传播编译区域之后，这是由于 AOTAutograd 编译函数与 dispatcher hook 的交互方式。

优化 Dynamo 与 DDP 的基本策略在 distributed.py 中进行了概述，其主要思想是在 DDP bucket 边界 上进行图中断。

当 DDP 中的每个节点需要与其节点同步权重时，它会将梯度和参数组织成 buckets，从而减少通信时间，并允许节点广播其一部分梯度给其他等待的节点。

分布式代码中的图中断意味着您可以期望 Dynamo 及其后端优化分布式程序的计算开销，但无法优化其通信开销。图中断可能会影响编译加速，如果减小的图大小会剥夺编译器融合的机会。然而，随着图大小的增加，收益会递减，因为目前大多数计算优化是局部融合。因此，在实践中，这种方法可能是足够的。

我还需要导出整个图吗？

对于绝大多数模型，您可能不需要，并且可以按原样使用 `torch.compile()`，但有几种情况需要完整的图，您可以通过简单地运行 `torch.compile(..., fullgraph=True)` 来确保一个完整的图。这些情况包括：

• 大规模训练运行，例如 $250K+，需要流水线并行和其他高级分片策略。

• 推理优化器，例如 TensorRT 或 AITemplate，它们依赖于比训练优化器更激进的融合。

• 移动端训练或推理。

未来的工作将包括将通信操作跟踪到图中，将这些操作与计算优化协调，以及优化通信操作。

为什么我的代码会崩溃？

如果您的代码在没有 `torch.compile` 的情况下运行正常，但在启用它后开始崩溃，那么最重要的第一步是弄清楚您的故障发生在堆栈的哪个部分。要解决此问题，请按照以下步骤操作，并且只有在前一个步骤成功后才尝试下一个步骤。

1. `torch.compile(..., backend="eager")` 仅运行 TorchDynamo 前向图捕获，然后使用 PyTorch 运行捕获的图。如果这失败了，那么 TorchDynamo 就存在问题。

2. `torch.compile(..., backend="aot_eager")` 运行 TorchDynamo 捕获前向图，然后 AOTAutograd 在没有任何额外的后端编译器步骤的情况下跟踪后向图。然后 PyTorch eager 将用于运行前向和后向图。如果这失败了，那么 AOTAutograd 就存在问题。

3. `torch.compile(..., backend="inductor")` 运行 TorchDynamo 捕获前向图，然后 AOTAutograd 使用 TorchInductor 编译器跟踪后向图。如果这失败了，那么 TorchInductor 就存在问题。

为什么编译很慢？

• Dynamo 编译 – TorchDynamo 内置了一个统计函数，用于收集和显示每个编译阶段花费的时间。这些统计数据可以通过在执行 `torch._dynamo` 后调用 `torch._dynamo.utils.compile_times()` 来访问。默认情况下，这将返回按名称计费的每个 TorchDynamo 函数的编译时间字符串表示。

• Inductor 编译– TorchInductor 内置了统计和跟踪功能，用于显示每个编译阶段花费的时间、输出代码、输出图可视化和 IR 转储。通过运行 `env TORCH_COMPILE_DEBUG=1 python repro.py` 来启用。这是一个旨在方便调试/理解 TorchInductor 内部机制的调试工具，输出内容大致如下所示：此链接。此调试跟踪中的每个文件都可以通过 `torch._inductor.config.trace.*` 进行启用/禁用。配置文件和图表默认都已禁用，因为生成它们成本较高。有关更多示例，请参见示例调试目录输出。

• 过度重新编译 当 TorchDynamo 编译一个函数（或其一部分）时，它会对局部变量和全局变量做出某些假设，以允许编译器进行优化，并将这些假设表示为在运行时检查特定值的 guard。如果任何一个 guard 失败，Dynamo 将会重新编译该函数（或其一部分），最多可重新编译 `torch._dynamo.config.recompile_limit` 次。如果您的程序达到了缓存限制，您首先需要确定是哪个 guard 失败以及是您程序的哪个部分触发了它。使用 `TORCH_TRACE/tlparse` 或 `TORCH_LOGS=recompiles` 来追溯问题的根源，有关更多详细信息，请参阅torch.compile 故障排除。

为什么在生产环境中重新编译？

在某些情况下，您可能不希望程序预热后出现意外的编译。例如，如果您在一个延迟敏感的应用程序中提供生产流量。为此，TorchDynamo 提供了一种替代模式，在该模式下使用之前编译的图，但不生成新的图。

frozen_toy_example = dynamo.run(toy_example)
frozen_toy_example(torch.randn(10), torch.randn(10))

您是如何加速我的代码的？

加速 PyTorch 代码有 3 种主要方式：

1. 通过垂直融合进行核函数融合，它融合了顺序操作以避免过多的读写。例如，融合两个连续的余弦函数意味着您可以进行 1 次读写，而不是 2 次读写 2 次写 2 次读。水平融合：最简单的例子是批处理，其中一个矩阵与一批示例相乘，但更通用的情况是分组 GEMM，其中一组矩阵乘法一起调度。

2. 无序执行：编译器的通用优化，通过向前查看图中的确切数据依赖关系，我们可以决定执行节点的最合适时机以及哪些缓冲区可以重用。

3. 自动工作放置：类似于无序执行点，但通过将图的节点与物理硬件或内存等资源匹配，我们可以设计一个合适的调度。

以上是加速 PyTorch 代码的通用原则，但不同的后端在优化方面会做出不同的权衡。例如，Inductor 首先负责融合它能融合的所有内容，然后才生成Triton核函数。

Triton 还通过自动内存合并、内存管理和流式多处理器内部调度来提供加速，并且它被设计用来处理分块计算。

然而，无论您使用哪个后端，最好使用基准测试和观察的方法，所以尝试使用 PyTorch 性能分析器，直观地检查生成的核函数，并尝试自己找出问题所在。

为什么我看不到加速效果？

图中断

通过 dynamo 看不到预期的加速的主要原因是图中断过多。那么什么是图中断呢？

给定一个程序，例如

def some_fun(x):
    ...

torch.compile(some_fun)(x)
...

TorchDynamo 将尝试将 `some_fun()` 中的所有 torch/tensor 操作编译成一个单独的 FX 图，但它可能无法将所有内容捕获到一个图中。

某些图中断的原因对于 TorchDynamo 来说是无法克服的，例如调用 PyTorch 以外的 C 扩展，TorchDynamo 无法看到这些扩展，并且可能会执行任意操作，而 TorchDynamo 无法引入必要的 guard 来确保编译后的程序可以安全重用。

为了最大化性能，尽可能减少图中断至关重要。

识别图中断的原因

要识别程序中的所有图中断及其原因，可以使用 `torch._dynamo.explain`。此工具会在提供的函数上运行 TorchDynamo，并聚合遇到的图中断。以下是示例用法：

import torch
import torch._dynamo as dynamo
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    print("woo")
    if b.sum() < 0:
        b = b * -1
    return x * b
explanation = dynamo.explain(toy_example)(torch.randn(10), torch.randn(10))
print(explanation)
"""
Graph Count: 3
Graph Break Count: 2
Op Count: 5
Break Reasons:
  Break Reason 1:
    Reason: builtin: print [<class 'torch._dynamo.variables.constant.ConstantVariable'>] False
    User Stack:
      <FrameSummary file foo.py, line 5 in toy_example>
  Break Reason 2:
    Reason: generic_jump TensorVariable()
    User Stack:
      <FrameSummary file foo.py, line 6 in torch_dynamo_resume_in_toy_example_at_5>
Ops per Graph:
  ...
Out Guards:
  ...
"""

要为遇到的第一个图中断抛出错误，您可以通过使用 `fullgraph=True` 来禁用 Python 回退，如果您曾使用过基于导出（export-based）的编译器，这对您来说应该很熟悉。

def toy_example(a, b):
   ...

torch.compile(toy_example, fullgraph=True, backend=<compiler>)(a, b)

为什么在我更改代码后没有重新编译？

如果您通过设置 `env TORCHDYNAMO_DYNAMIC_SHAPES=1 python model.py` 启用了动态形状，那么您的代码在形状更改时不会重新编译。我们增加了对动态形状的支持，这可以避免形状变化小于 2 倍时的重新编译。这在计算机视觉中的可变图像大小或自然语言处理中的可变序列长度等场景中特别有用。在推理场景中，由于您接收的是来自不同客户端应用程序的批次，因此通常无法提前知道批次大小。 2022 年 10 月 25 日

总的来说，TorchDynamo 会尽力避免不必要的重新编译，因此，例如，如果 TorchDynamo 找到 3 个图而您的更改只修改了一个图，那么只有那个图会重新编译。因此，另一个避免潜在的慢编译时间的技巧是预热模型，即先编译一次，然后后续编译会快得多。冷启动编译时间仍然是我们可见跟踪的指标。

为什么我得到的结果不正确？

通过设置环境变量 `TORCHDYNAMO_REPRO_LEVEL=4` 还可以最大限度地减少准确性问题，它以类似 git bisect 的模型运行，并且一个完整的重现可能类似于 `TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4`。我们需要这个的原因是下游编译器会生成代码，无论是 Triton 代码还是 C++ 后端，这些下游编译器的数值可能存在细微差异，但会对您的训练稳定性产生巨大的影响。因此，准确性调试器对我们检测其代码生成或后端编译器中的错误非常有用。

如果您想确保随机数生成在 torch 和 triton 之间保持一致，您可以启用 `torch._inductor.config.fallback_random = True`。

为什么我遇到了 OOM（内存不足）？

Dynamo 仍处于 alpha 阶段，因此有一些 OOM 的来源，如果您遇到 OOM，请尝试按顺序禁用以下配置，然后打开一个 GitHub 问题，以便我们解决根本问题：1. 如果您正在使用动态形状，请尝试禁用它们，我们默认禁用了它们：`env TORCHDYNAMO_DYNAMIC_SHAPES=0 python model.py` 2. Triton 的 CUDA 图默认在 Inductor 中启用，但移除它们可能会缓解一些 OOM 问题：`torch._inductor.config.triton.cudagraphs = False`。

`torch.func` 是否与 `torch.compile` 一起使用（用于 `grad` 和 `vmap` 变换）？

将 `torch.func` 变换应用于使用 `torch.compile` 的函数是有效的。

import torch

@torch.compile
def f(x):
    return torch.sin(x)

def g(x):
    return torch.grad(f)(x)

x = torch.randn(2, 3)
g(x)

在由 `torch.compile` 处理的函数内部调用 `torch.func` 变换

使用 `torch.compile` 编译 `torch.func.grad`

import torch

def wrapper_fn(x):
    return torch.func.grad(lambda x: x.sin().sum())(x)

x = torch.randn(3, 3, 3)
grad_x = torch.compile(wrapper_fn)(x)

使用 `torch.compile` 编译 `torch.vmap`

import torch

def my_fn(x):
    return torch.vmap(lambda x: x.sum(1))(x)

x = torch.randn(3, 3, 3)
output = torch.compile(my_fn)(x)

编译不支持的函数（逃生舱）

对于其他变换，作为一种变通方法，请使用 `torch._dynamo.allow_in_graph`。

`allow_in_graph` 是一个逃生舱。如果您的代码无法与检查 Python 字节码的 `torch.compile` 配合使用，但您认为它可以与符号跟踪方法（如 `jax.jit`）配合使用，那么请使用 `allow_in_graph`。

通过使用 `allow_in_graph` 注释一个函数，您必须确保您的代码满足以下要求：

• 您的函数中的所有输出仅依赖于输入，而不依赖于任何捕获的 Tensor。

• 您的函数是函数式的。也就是说，它不会修改任何状态。这可能会放宽；我们实际上支持从外部看起来是函数式的函数：它们可能有就地（in-place）的 PyTorch 操作，但不得修改全局状态或函数的输入。

• 您的函数不会引发数据相关的错误。

import torch

@torch.compile
def f(x):
    return torch._dynamo.allow_in_graph(torch.vmap(torch.sum))(x)

x = torch.randn(2, 3)
f(x)

一个常见的陷阱是使用 `allow_in_graph` 注释调用 `nn.Module` 的函数。这是因为输出现在依赖于 `nn.Module` 的参数。要使其正常工作，请使用 `torch.func.functional_call` 来提取模块状态。

`torch.compile` 是否支持 NumPy？

从 2.1 版本开始，`torch.compile` 理解操作 NumPy 数组的原生 NumPy 程序，以及通过 `x.numpy()`、`torch.from_numpy` 和相关函数从 PyTorch 转换为 NumPy 再转换回来的混合 PyTorch-NumPy 程序。

`torch.compile` 支持哪些 NumPy 功能？

NumPy 在 `torch.compile` 中遵循 NumPy 2.0 预发布版。

通常，`torch.compile` 能够跟踪大多数 NumPy 构造，当它无法跟踪时，它会回退到即时执行，并让 NumPy 执行那段代码。即使如此，也有一些功能是 `torch.compile` 的语义与 NumPy 的语义略有偏差。

• NumPy 标量：我们将其建模为 0 维数组。也就是说，在 `torch.compile` 下，`np.float32(3)` 返回一个 0 维数组。为了避免图中断，最好使用此 0 维数组。如果这会中断您的代码，您可以通过将 NumPy 标量转换为相关的 Python 标量类型（`bool/int/float`）来解决。

• 负步幅：`np.flip` 和具有负步幅的切片会返回副本。

• 类型提升：NumPy 的类型提升规则将在 NumPy 2.0 中更改。新规则在NEP 50 中进行了描述。`torch.compile` 实现了 NEP 50，而不是当前即将被弃用的规则。

• `{tril,triu}_indices_from/{tril,triu}_indices` 返回数组而不是数组的元组。

还有其他我们不支持跟踪的功能，我们会优雅地回退到 NumPy 来执行它们。

• 非数字 dtype，如日期时间、字符串、字符、void、结构化 dtype 和 recarrays。

• 长 dtype `np.float128/np.complex256` 和一些无符号 dtype `np.uint16/np.uint32/np.uint64`。

• `ndarray` 子类。

• 掩码数组。

• 不常见的 ufunc 机制，如 `axes=[(n,k),(k,m)->(n,m)]` 和 ufunc 方法（例如，`np.add.reduce`）。

• 对 `complex64/complex128` 数组的排序/排序。

• NumPy `np.poly1d` 和 `np.polynomial`。

• 具有 2 个或更多返回值的函数中的位置参数 `out1, out2`（`out=tuple` 可以正常工作）。

• `__array_function__`, `__array_interface__` 和 `__array_wrap__`。

• `ndarray.ctypes` 属性。

我可以使用 `torch.compile` 编译 NumPy 代码吗？

当然可以！`torch.compile` 原生支持 NumPy 代码，并将其视为 PyTorch 代码。为此，只需使用 `torch.compile` 装饰器包装 NumPy 代码即可。

import torch
import numpy as np

@torch.compile
def numpy_fn(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = np.random.randn(1024, 64)
Y = np.random.randn(1024, 64)
Z = numpy_fn(X, Y)
assert isinstance(Z, np.ndarray)

使用环境变量 `TORCH_LOGS=output_code` 执行此示例，我们可以看到 `torch.compile` 能够将乘法和求和融合到一个 C++ 核函数中。它还可以使用 OpenMP（原生 NumPy 是单线程的）并行执行它们。这可以轻松地使您的 NumPy 代码快 n 倍，其中 n 是您处理器核心的数量！

通过这种方式跟踪 NumPy 代码也支持编译代码中的图中断。

我可以在 CUDA 上执行 NumPy 代码并通过 `torch.compile` 计算梯度吗？

是的，您可以！要做到这一点，您只需在 `torch.device("cuda")` 上下文中执行您的代码。考虑以下示例：

import torch
import numpy as np

@torch.compile
def numpy_fn(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = np.random.randn(1024, 64)
Y = np.random.randn(1024, 64)
with torch.device("cuda"):
    Z = numpy_fn(X, Y)
assert isinstance(Z, np.ndarray)

在此示例中，`numpy_fn` 将在 CUDA 上执行。为此，`torch.compile` 会自动将 `X` 和 `Y` 从 CPU 移动到 CUDA，然后它会将结果 `Z` 从 CUDA 移动到 CPU。如果我们是在同一个程序运行中多次执行此函数，我们可能希望避免所有这些相对昂贵的内存复制。为此，我们只需调整我们的 `numpy_fn`，使其接受 cuda Tensor 并返回 Tensor。我们可以通过使用 `torch.compiler.wrap_numpy` 来实现。

@torch.compile(fullgraph=True)
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = torch.randn(1024, 64, device="cuda")
Y = torch.randn(1024, 64, device="cuda")
Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
assert Z.device.type == "cuda"

在这里，我们显式地在 CUDA 内存中创建 Tensor，并将它们传递给函数，该函数在 CUDA 设备上执行所有计算。`wrap_numpy` 负责在 `torch.compile` 级别将任何 `torch.Tensor` 输入标记为具有 `np.ndarray` 语义的输入。在编译器内部标记 Tensor 是一项非常便宜的操作，因此在运行时不会发生数据复制或数据移动。

使用此装饰器，我们还可以通过 NumPy 代码进行微分！

@torch.compile(fullgraph=True)
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    return np.mean(np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1)))

X = torch.randn(1024, 64, device="cuda", requires_grad=True)
Y = torch.randn(1024, 64, device="cuda")
Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
Z.backward()
# X.grad now holds the gradient of the computation
print(X.grad)

我们一直在使用 `fullgraph=True`，因为在这种情况下图中断是有问题的。当发生图中断时，我们需要具体化 NumPy 数组。由于 NumPy 数组没有 `device` 或 `requires_grad` 的概念，此信息在图中断期间会丢失。

我们无法将梯度传播通过图中断，因为图中断代码可能执行任意无法区分的代码。另一方面，在 CUDA 执行的情况下，我们可以通过使用 `torch.device("cuda")` 上下文管理器来解决此问题，正如我们在第一个示例中所做的那样。

@torch.compile
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    prod = X[:, :, None] * Y[:, None, :]
    print("oops, a graph break!")
    return np.sum(prod, axis=(-2, -1))

X = torch.randn(1024, 64, device="cuda")
Y = torch.randn(1024, 64, device="cuda")

with torch.device("cuda"):
    Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
assert Z.device.type == "cuda"

在图中断期间，中间 Tensor 仍然需要移动到 CPU，但当在图中断后恢复跟踪时，图的其余部分仍在 CUDA 上进行跟踪。考虑到这种 CUDA <> CPU 和 CPU <> CUDA 的移动，图中断在 NumPy 上下文中非常昂贵，应避免，但至少它们允许跟踪复杂的代码片段。

如何调试 `torch.compile` 下的 NumPy 代码？

调试 JIT 编译的代码具有挑战性，因为现代编译器的复杂性和它们引发的令人望而生畏的错误。torch.compile 故障排除文档 包含一些解决此任务的技巧。

如果以上内容不足以确定问题的根源，我们仍然可以使用其他一些特定于 NumPy 的工具。我们可以通过禁用对 NumPy 函数的跟踪来区分错误是否完全在 PyTorch 代码中。

from torch._dynamo import config
config.trace_numpy = False

如果错误存在于跟踪的 NumPy 代码中，我们可以使用 PyTorch 作为后端，通过导入 `import torch._numpy as np` 来即时执行 NumPy 代码（不使用 `torch.compile`）。这应该仅用于**调试目的**，绝不能替代 PyTorch API，因为它**性能差得多**，而且作为私有 API，**可能会在不通知的情况下更改**。无论如何，`torch._numpy` 是一个基于 PyTorch 的 NumPy 的 Python 实现，它由 `torch.compile` 在内部用于将 NumPy 代码转换为 PyTorch 代码。它相对容易阅读和修改，因此如果您发现任何错误，请随时提交 PR 进行修复或简单地打开一个问题。

如果在导入 `torch._numpy as np` 时程序可以工作，那么很可能是 TorchDynamo 中存在错误。如果是这种情况，请随意打开一个带有最小重现器的问题。

我使用 `torch.compile` 编译了一些 NumPy 代码，但没有看到任何加速。

开始的最佳地点是关于如何调试此类 torch.compile 问题的通用建议的教程。

由于使用了不支持的功能，可能会发生一些图中断。请参阅torch.compile 支持哪些 NumPy 功能？。更普遍地说，要记住某些广泛使用的 NumPy 功能与编译器不兼容。例如，就地修改使得编译器内部的推理变得困难，并且通常会产生比其原地对应物更差的性能。因此，最好避免使用它们。使用 `out=` 参数也是如此。相反，请优先使用原地操作，让 `torch.compile` 优化内存使用。数据相关操作（如布尔掩码的掩码索引）或数据相关的控制流（如 `if` 或 `while` 结构）也是如此。

使用哪个 API 进行细粒度跟踪？

在某些情况下，您可能需要将代码的小部分排除在 `torch.compile` 编译之外。本节提供了一些答案，您可以在用于细粒度跟踪的 TorchDynamo API 中找到更多信息。

如何对函数进行图中断？

对函数进行图中断不足以充分表达您希望 PyTorch 执行的操作。您需要更具体地说明您的用例。一些最常见的用例，您可能需要考虑：

• 如果您想禁用此函数帧及其递归调用的帧的编译，请使用 `torch._dynamo.disable`。

• 如果您希望特定运算符（如 `fbgemm`）使用即时模式，请使用 `torch._dynamo.disallow_in_graph`。

一些不常见的用例包括：

• 如果您想禁用函数帧上的 TorchDynamo，但在递归调用的帧上重新启用它 – 请使用 `torch._dynamo.disable(recursive=False)`。

• 如果您想阻止内联函数帧 – 请在您想阻止内联的函数开头使用 `torch._dynamo.graph_break`。

`torch._dynamo.disable` 和 `torch._dynamo.disallow_in_graph` 有什么区别？

Disallow-in-graph 在运算符级别工作，更具体地说，在 TorchDynamo 提取的图中所见的运算符级别工作。

Disable 在函数帧级别工作，并决定 TorchDynamo 是否应查看函数帧。

`torch._dynamo.disable` 和 `torch._dynamo_skip` 有什么区别？

注意

`torch._dynamo_skip` 已弃用。

您最有可能需要 `torch._dynamo.disable`。但在不太可能的情况下，您可能需要更精细地控制。假设您只想禁用 `a_fn` 函数上的跟踪，但希望在 `aa_fn` 和 `ab_fn` 中继续跟踪。下图演示了此用例。

在这种情况下，您可以使用 `torch._dynamo.disable(recursive=False)`。在以前的版本中，此功能由 `torch._dynamo.skip` 提供。现在由 `torch._dynamo.disable` 中的 `recursive` 标志支持。