常见问题

图表中断

你使用 Dynamo 没有看到预期的加速效果，主要原因是过多的图断开。那么，什么是图断开呢？

def some_fun(x):
    ...

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

Torchdynamo 尝试将 some_fun() 中的所有 torch/tensor 操作编译成一个单一的 FX 图，但可能会失败，无法将所有操作都包含在一个图中。

一些图中断的原因对于TorchDynamo来说是无法克服的，例如调用除PyTorch以外的C扩展对TorchDynamo来说是不可见的。这些扩展可能在没有给TorchDynamo引入必要的保护措施的情况下执行任意操作，导致编译后的程序不安全地重用。

为了最大化性能，应该尽可能减少图表中的断点。

确定图中断的原因

要识别程序中所有图中断及其原因，可以使用 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 回退功能。如果你之前使用过基于导出的编译器，这一点应该很熟悉。

def toy_example(a, b):
   ...

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

为什么我修改代码后没有触发重新编译？

如果你通过设置 env TORCHDYNAMO_DYNAMIC_SHAPES=1 python model.py 启用了动态形状功能，那么当输入的形状发生变化时，你的代码不会重新编译。我们添加了对动态形状的支持，在这种情况下，如果形状的变化小于两倍，则可以避免重新编译。这在诸如计算机视觉中图像大小变化或自然语言处理中序列长度可变的情况下特别有用。在推理场景中，由于你从不同的客户端应用程序获取数据，通常无法提前确定批次的大小。

通常，TorchDynamo 努力避免不必要的重新编译。例如，如果 TorchDynamo 找到 3 个图而你的更改只修改了一个图，则只会重新编译该图。因此，另一个技巧是先预热模型（即编译一次），这样后续的编译会快得多。冷启动编译时间是我们显式跟踪的一个指标。

在使用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 是一个逃生机制。如果你的代码与 torch.compile 不兼容，而 torch.compile 会分析 Python 字节码，但你认为通过符号跟踪方法（如 jax.jit）可以解决问题，则使用 allow_in_graph。

通过在函数上使用allow_in_graph注解，你必须确保你的代码满足以下要求：

• 你的函数中所有的输出仅取决于输入，不依赖于任何被捕获的张量。

• 你的函数是纯函数式的，即不修改任何状态。这一要求可以适当放宽；我们实际支持的是那些从外部看起来像纯函数的函数：它们可以包含就地执行的 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)

一个常见的陷阱是在调用 nn.Module 的函数上使用 allow_in_graph 进行注解。这是因为输出现在依赖于 nn.Module 的参数。为了使这能正常工作，应使用 torch.func.functional_call 来提取模块状态。

torch.compile 支持哪些 NumPy 功能？

在torch.compile中使用的NumPy遵循NumPy 2.0的预发布版本。

通常，torch.compile 可以追踪大多数的 NumPy 构造。当无法追踪时，它会退回到即时执行模式，并让 NumPy 执行那段代码。即使如此，仍然有一些特性使得 torch.compile 的行为与 NumPy 稍微不同：

• NumPy 标量：我们将它们建模为0维数组。例如，np.float32(3) 在 torch.compile 下返回一个 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 进行执行。

• 如日期时间、字符串、字符、void、结构化数据类型和 recarrays 等非数值数据类型。

• 长整型数据类型 np.float128/np.complex256 以及一些无符号数据类型 np.uint16/np.uint32/np.uint64。

• ndarray 子类。

• 带掩码的数组。

• 类似于axes=[(n,k),(k,m)->(n,m)]的esoteric ufunc机制以及ufunc方法（例如np.add.reduce）。

• 对 complex64/complex128 数组进行排序。

• NumPy 的 np.poly1d 和 np.polynomial。

• 当函数返回两个或更多的值时，可以使用位置参数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 再次移动回 CPU。如果我们在这个程序运行中多次执行此函数，我们可能希望避免这些昂贵的内存复制操作。为此，我们可以调整 numpy_fn 使其接受和返回 cuda 张量，并使用 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内存中创建张量，并将它们传递给函数，在该函数中所有计算都在CUDA设备上完成。wrap_numpy 负责标记任何 torch.Tensor 输入为具有 np.ndarray 语义的输入，这是在 torch.compile 层级上进行的。在编译器内部对张量进行标记是一个非常便宜的操作，因此在运行时不会发生数据复制或移动。

使用这个装饰器，我们可以对 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"

在图中断期间，中间张量仍然需要移动到CPU上。但在恢复追踪后，剩余的图部分仍在CUDA上进行追踪。由于这种CUDA与CPU之间的频繁切换，在NumPy上下文中，图中断的成本较高，应尽量避免。不过，至少它们允许通过复杂的代码段进行追踪。

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

调试即时编译的代码颇具挑战性，考虑到现代编译器的复杂性和它们引发的令人难以处理的错误。如何在 torch.compile 中诊断运行时错误的教程提供了一些应对这一任务的小技巧。

如果上述方法仍无法确定问题的来源，我们还可以使用一些特定于NumPy的工具进行进一步排查。通过禁用NumPy函数中的跟踪功能，我们可以判断错误是否完全由PyTorch代码引起。

from torch._dynamo import config
config.trace_numpy = False

如果错误出现在被追踪的NumPy代码中，我们可以使用PyTorch作为后端（不使用torch.compile），通过导入import torch._numpy as np来急切地执行NumPy代码。这仅用于调试目的，绝不是PyTorch API的替代品，因为它的性能要差得多，并且作为私有API，可能会在没有通知的情况下更改。无论如何，torch._numpy是基于PyTorch实现的NumPy的Python版本，它被torch.compile内部使用来将NumPy代码转换为PyTorch代码。它的读取和修改都非常容易，因此如果你发现其中有任何错误，请随时提交修复该问题的PR或直接打开一个issue。

如果导入 torch._numpy as np 后程序可以正常运行，那么问题可能出在 TorchDynamo 中。如果是这种情况，请提交一个包含最小复现代码的问题。

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

最佳起点是参考这个教程，它提供了关于如何调试此类 torch.compile 问题的一般性建议。

一些图中断可能是由于使用了不受支持的功能导致的。参见torch.compile 支持哪些 NumPy 功能？。通常需要注意的是，某些广泛使用的 NumPy 功能与编译器不兼容。例如，就地修改使得在编译器中进行推理变得困难，并且性能往往不如非就地操作好。out= 参数同样需要避免使用。建议使用非就地操作，并让 torch.compile 优化内存使用。此外，数据依赖的操作（如通过布尔掩码进行的索引）和数据依赖的控制流（例如 if 或 while 构造）也需要特别注意。

如何在函数上设置断点？

仅仅在函数上设置断点是不够的，无法充分表达你希望 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.disable。但在极少数情况下，你可能需要更细粒度的控制。假设你想只在 a_fn 函数中禁用跟踪，但希望继续在 aa_fn 和 ab_fn 中进行跟踪。下面的图片演示了这种情况：

在这种情况下，你可以使用 torch._dynamo.disable(recursive=False)。在之前的版本中，这项功能是由 torch._dynamo.skip 提供的。现在可以通过设置 torch._dynamo.disable 中的 recursive 标志来实现。