常见问题解答¶
作者: Mark Saroufim
torch.compile
支持训练吗?¶
torch.compile
支持训练,使用 AOTAutograd 来捕获反向传播
The
.forward()
图和optimizer.step()
被 TorchDynamo 的 pythonevalframe
前端捕获。对于 TorchDynamo 捕获的每个
.forward()
段,它使用 AOTAutograd 来生成反向传播图段。每对正向和反向图(可选)进行最小割分区,以保存正向和反向之间的最小状态。
正向和反向对被包装在
autograd.function
模块中。用户代码调用
.backward()
仍然会触发 eager 的 autograd 引擎,它会将每个编译后的反向传播图视为一个运算符来运行,还会运行任何非编译 eager 运算符的.backward()
函数。
您支持分布式代码吗?¶
torch.compile
支持 DistributedDataParallel
(DDP)。正在考虑支持其他分布式训练库。
分布式代码在 dynamo 中具有挑战性的主要原因是 AOTAutograd 会展开正向和反向传播,并为后端提供 2 个图进行优化。这对分布式代码来说是一个问题,因为我们希望理想情况下将通信操作与计算操作重叠。Eager PyTorch 使用 autograd 钩子、模块钩子和模块状态的修改/变异来以不同的方式实现这一点,用于 DDP/FSDP。在 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 并将您的程序在仅观察的“编译器”下运行的过程,该编译器记录任何保护失败的原因。 您应该确保运行程序的时间(迭代次数)至少与遇到问题时运行的时间一样长,分析器将在此期间累积统计信息。
from torch._dynamo.utils import CompileProfiler
def my_model():
...
with CompileProfiler() as prof:
profiler_model = torch.compile(my_model, backend=prof)
profiler_model()
print(prof.report())
为什么您在生产环境中重新编译?¶
在某些情况下,您可能不希望程序预热后出现意外编译。 例如,如果您正在延迟敏感的应用程序中提供生产流量。 为此,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 分析器,可视化检查生成的内核,并尝试自己了解发生了什么。
为什么我没有看到加速效果?¶
图中断¶
使用 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
启用了动态形状,那么您的代码在形状发生更改时将不会重新编译。 我们添加了对动态形状的支持,这避免了在形状变化小于 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. 带有 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
是一个应急措施。如果您的代码无法与 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 结构,而当它无法跟踪时,它会回退到急切模式,并让 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
。带有 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 张量并返回张量。我们可以通过使用 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 作为后端来急切地执行 NumPy 代码(不使用 torch.compile
),方法是导入 import torch._numpy as np
。这应该只用于 **调试目的**,它绝不是 PyTorch API 的替代品,因为它 **效率低得多**,而且作为私有 API,**可能会在未经通知的情况下发生更改**。无论如何, torch._numpy
是 NumPy 的一个 Python 实现,它使用 PyTorch,并由 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
之间的区别是什么?¶
禁止在图中工作在运算符级别,更具体地说,是在 TorchDynamo 提取的图中看到的运算符级别。
禁用在函数帧级别工作,并决定 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
标志支持。