常见问题解答¶

`torch.compile` 支持训练吗？¶

torch.compile 支持训练，使用 AOTAutograd 捕获反向

.forward() 图和 optimizer.step() 由 TorchDynamo 的 python evalframe 前端捕获。
对于 torchdynamo 捕获的每个 .forward() 片段，它使用 AOTAutograd 生成一个反向图片段。
每对前向和反向图都（可选地）进行最小割分区，以保存前向和反向之间的最小状态。
前向和反向对被包装在 autograd.function 模块中。
用户代码调用 .backward() 仍然会触发 eager 的 autograd 引擎，该引擎将每个 *编译后的反向* 图作为单个操作运行，同时还运行任何未编译的 eager 操作的 .backward() 函数。

您是否支持分布式代码？¶

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

Dynamo 难以处理分布式代码的主要原因是，AOTAutograd 会展开前向和反向传递，并为后端提供 2 个图进行优化。这对分布式代码来说是一个问题，因为我们希望理想情况下将通信操作与计算操作重叠。Eager PyTorch 通过不同的方式（对于 DDP/FSDP）实现这一点 - 使用 autograd 钩子、模块钩子和模块状态的修改/变异。在 Dynamo 的简单应用中，由于 AOTAutograd 编译函数与调度程序钩子的交互方式，在反向传播期间应该在操作之后立即运行的钩子可能会延迟到整个编译区域的反向传播操作之后。

使用 Dynamo 优化 DDP 的基本策略在 distributed.py 中概述，其主要思想是在 DDP 桶边界上进行图分割。

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

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

我是否仍然需要导出整个图？¶

对于绝大多数模型，您可能不需要，您可以按原样使用 torch.compile()，但在某些情况下，完整图是必要的，您可以通过简单地运行 torch.compile(..., fullgraph=True) 来确保完整图。这些情况包括

大规模训练运行，例如需要流水线并行和其他高级分片策略的 250K+。
像 TensorRT 或 AITemplate 这样的推理优化器，它们依赖于比训练优化器更积极的融合。
移动训练或推理。

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

为什么我的代码崩溃了？¶

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

torch.compile(..., backend="eager") 仅运行 TorchDynamo 前向图捕获，然后使用 PyTorch 运行捕获的图。如果这失败了，那么 TorchDynamo 存在问题。
torch.compile(..., backend="aot_eager") 会运行 TorchDynamo 捕获前向图，然后使用 AOTAutograd 追踪反向图，而无需任何额外的后端编译步骤。然后 PyTorch Eager 将用于运行前向和反向图。如果这失败了，则说明 AOTAutograd 存在问题。
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 编译一个函数（或函数的一部分）时，它会对局部变量和全局变量做出某些假设，以便允许编译器优化，并将这些假设表示为在运行时检查特定值的保护。如果任何一个保护失败，Dynamo 将重新编译该函数（或部分）最多 torch._dynamo.config.cache_size_limit 次。如果你的程序达到了缓存限制，你需要先确定哪个保护失败以及程序的哪个部分触发了它。重新编译分析器自动执行将 TorchDynamo 的缓存限制设置为 1 并在你只观察的“编译器”下运行程序的过程，该编译器会记录任何保护失败的原因。你应该确保你的程序至少运行与遇到问题时一样长的时间（迭代次数），分析器将在此期间累积统计数据。

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

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

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

如何加速我的代码？¶

有三种主要方法可以加速 PyTorch 代码。

通过垂直融合进行内核融合，将顺序操作融合以避免过多的读/写。例如，融合两个连续的余弦函数意味着你可以执行 1 次读 1 次写，而不是 2 次读 2 次写 2 次。水平融合：最简单的例子是批处理，其中一个矩阵与一批示例相乘，但更一般的场景是分组 GEMM，其中一组矩阵乘法被一起调度。
乱序执行：编译器的一般优化，通过前瞻性地查看图中确切的数据依赖关系，我们可以确定执行节点的最合适时间以及哪些缓冲区可以重用。
自动工作放置：类似于乱序执行点，但通过将图的节点与物理硬件或内存等资源匹配，我们可以设计一个合适的调度方案。

以上是加速 PyTorch 代码的一般原则，但不同的后端将在优化内容上做出不同的权衡。例如，Inductor 首先会尽其所能进行融合，然后才会生成 Triton 内核。

此外，Triton 还提供了加速，因为它可以自动合并内存、管理内存并在每个流式多处理器中进行调度，并且它被设计用于处理平铺计算。

但是，无论你使用哪个后端，最好使用基准测试并查看方法，因此请尝试使用 PyTorch 分析器，直观地检查生成的内核并尝试自己了解发生了什么。

为什么我没有看到加速？¶

图中断¶

你没有看到期望的加速的主要原因是图中断过多。那么什么是图中断呢？

给定一个类似于以下的程序：

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 启用了动态形状，那么你的代码在形状发生变化时不会重新编译。我们添加了对动态形状的支持，这避免了在形状变化小于 2 倍的情况下重新编译。这在 CV 中图像大小变化或 NLP 中序列长度变化等场景中特别有用。在推理场景中，通常无法预先知道批处理大小是多少，因为你从不同的客户端应用程序获取你所能获取的内容。

通常，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. CUDA 图与 Triton 在 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 是一个逃生舱。如果你的代码不适用于 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)

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

NumPy 是否与`torch.compile`一起使用？¶

从 2.1 版本开始，torch.compile 能够理解在 NumPy 数组上工作的原生 NumPy 程序，以及通过x.numpy()、torch.from_numpy和相关函数在 PyTorch 和 NumPy 之间转换的混合 PyTorch-NumPy 程序。

`torch.compile`支持哪些 NumPy 特性？¶

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

通常，torch.compile能够跟踪大多数 NumPy 结构，当它无法跟踪时，它会回退到 Eager 模式并让 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 来执行它们。

非数值数据类型，如日期时间、字符串、字符、空值、结构化数据类型和记录数组。
长数据类型np.float128/np.complex256和一些无符号数据类型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。
在具有两个或多个返回值的函数中使用位置参数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 张量并返回张量。我们可以使用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.compile级别将任何torch.Tensor输入标记为具有np.ndarray语义的输入。在编译器内部标记张量是一个非常廉价的操作，因此在运行时不会发生任何数据复制或数据移动。

使用此装饰器，我们还可以通过 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 和 CPU <> CUDA 移动，图中断在 NumPy 上下文中相当昂贵，应避免，但至少它们允许跟踪复杂代码段。

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

鉴于现代编译器的复杂性和它们引发的令人望而生畏的错误，调试 JIT 编译的代码具有挑战性。有关如何在 torch.compile 中诊断运行时错误的教程包含一些关于如何解决此任务的提示和技巧。

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

from torch._dynamo import config
config.trace_numpy = False

如果错误存在于跟踪的 NumPy 代码中，我们可以使用 PyTorch 作为后端以 Eager 模式（无需torch.compile）执行 NumPy 代码，方法是导入import torch._numpy as np。这仅应用于**调试目的**，绝不是 PyTorch API 的替代品，因为它**性能低得多**，并且作为私有 API，**可能会在未经通知的情况下发生更改**。无论如何，torch._numpy是根据 PyTorch 使用 Python 实现的 NumPy，它由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 中继续跟踪。下图演示了此用例

diagram of torch.compile + disable(a_fn, recursive=False)

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

常见问题解答¶

torch.compile 支持训练吗？¶