PyTorch 2.0 故障排除

作者: Michael Lazos

我们正在积极开发调试工具、性能分析器，并改进我们的错误和警告信息。以下是可用工具及其典型用法的表格。如需更多帮助，请参阅诊断运行时错误。

Title

工具	目的	用法
信息日志记录	查看编译的汇总步骤	torch._logging.set_logs(dynamo = logging.INFO) 或 TORCH_LOGS="dynamo"
调试日志记录	查看编译的详细步骤（打印每条追踪的指令）	torch._logging.set_logs(dynamo = logging.DEBUG) 和 torch._dynamo.config.verbose = True，或者 TORCH_LOGS="+dynamo" TORCHDYNAMO_VERBOSE=1
适用于任何后端的压缩工具	找到最小的子图，该子图可以在任何后端重现错误	设置环境变量 TORCHDYNAMO_REPRO_AFTER="dynamo"
用于 TorchInductor 的压缩器	如果错误在 AOTAutograd 之后发生，找到在 TorchInductor 降低过程中重现错误的最小子图	设置环境变量 TORCHDYNAMO_REPRO_AFTER="aot"
Dynamo 精度最小化器	找到最小的子图，当您怀疑问题出在AOTAutograd时，该子图在急切模式模型和优化模型之间重现了精度问题。	TORCHDYNAMO_REPRO_AFTER="dynamo" TORCHDYNAMO_REPRO_LEVEL=4
电感精度缩小器	找到最小的子图，当您怀疑问题是出在后台（例如，inductor）时，该子图会在急切模式模型和优化模型之间重现精度问题。如果这不起作用，请尝试使用Dynamo精度缩小器。	TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4
torch._dynamo.explain	查找图形断点并显示其推理	torch._dynamo.explain(fn)(*inputs)
记录/回放	记录和重放帧以在图形捕获期间重现错误	torch._dynamo.config.replay_record_enabled = True
TorchDynamo 函数名称过滤	仅编译具有给定名称的函数，以减少调试问题时的噪音	设置环境变量 TORCHDYNAMO_DEBUG_FUNCTION=
TorchInductor 调试日志记录	打印一般 TorchInductor 调试信息和生成的 Triton/C++ 代码	torch._inductor.config.debug = True
TorchInductor 追踪	显示每个TorchInductor阶段所用的时间 + 输出代码和图形可视化	设置环境变量 TORCH_COMPILE_DEBUG=1 或 torch._inductor.config.trace.enabled = True

除了信息和调试日志记录外， 您还可以使用 torch._logging 进行更细粒度的日志记录。

诊断运行时错误

在高层次上，TorchDynamo 栈由从 Python 代码捕获的图（TorchDynamo）和后端编译器组成。例如，后端编译器可能包括反向图追踪（AOTAutograd）和图降低（TorchInductor）*。错误可能发生在栈的任何组件中，并将提供完整的栈追踪。

要确定错误发生在哪个组件中，您可以使用信息级别的日志记录 torch._logging.set_logs(dynamo = logging.INFO) 或 TORCH_LOGS="dynamo" 并查找 Step #: ... 输出。日志在每个步骤的开始和结束时生成，因此错误应对应的步骤是最近记录的步骤，其结束尚未记录。这些步骤对应于堆栈的以下部分：

步骤	组件
1	TorchDynamo
2	编译器后端
3	TorchInductor

如果信息日志记录不足，您可以使用可用的后端选项。这些选项包括：

• "eager": 仅运行 TorchDynamo 前向图捕获，然后使用 PyTorch 运行捕获的图。这提供了 TorchDynamo 是否引发错误的指示。

• "aot_eager": 运行 TorchDynamo 以捕获前向图，然后使用 AOTAutograd 跟踪后向图，无需任何额外的后端编译步骤。PyTorch eager 将用于运行前向和后向图。这对于将问题缩小到 AOTAutograd 非常有用。

缩小问题的通用步骤如下：

1. 使用"eager"后端运行您的程序。如果错误不再出现，问题在于所使用的后端编译器（如果使用TorchInductor，请继续执行步骤2。如果不是，请参阅此部分）。如果使用"eager"后端时错误仍然出现，则是在运行torchdynamo时发生的错误。

2. 此步骤仅在使用 TorchInductor 作为后端编译器时才需要。使用 "aot_eager" 后端运行模型。如果此后端引发错误，则错误发生在 AOTAutograd 跟踪期间。如果使用此后端不再发生错误，则 错误在 TorchInductor* 中。

这些案例将在以下章节中进行分析。

注意

TorchInductor 后端包括 AOTAutograd 跟踪和 TorchInductor 编译器本身。我们将通过将 TorchInductor 称为后端，并将 TorchInductor 降低称为降低由 AOTAutograd 跟踪的图的阶段来消除歧义。

Torchdynamo 错误

如果生成的错误发生在 "eager" 后端，那么 TorchDynamo 很可能是错误的来源。以下是一个会生成错误的示例代码。

import torch

import torch._dynamo as dynamo


def test_assertion_error():
    y = torch.ones(200, 200)
    z = {y: 5}
    return z

compiled_test_assertion_error = torch.compile(test_assertion_error, backend="eager")

compiled_test_assertion_error()

上述代码生成以下错误：

torch._dynamo.convert_frame: [ERROR] WON'T CONVERT test_assertion_error /scratch/mlazos/torchdynamo/../test/errors.py line 26
由于:
Traceback (most recent call last):
  File "/scratch/mlazos/torchdynamo/torchdynamo/symbolic_convert.py", line 837, in BUILD_MAP
    assert isinstance(k, ConstantVariable) or (
AssertionError

来自用户代码:
   File "/scratch/mlazos/torchdynamo/../test/errors.py", line 34, in test_assertion_error
    z = {y: 5}

设置 torch._dynamo.config.verbose=True 以获取更多信息
==========

正如消息所示，您可以设置 torch._dynamo.config.verbose=True 以获取完整的堆栈跟踪，包括 TorchDynamo 中的错误和用户代码。除了这个标志外，您还可以通过 torch._logging.set_logs(dynamo = logging.INFO) 或 TORCH_LOGS="dynamo" 来设置 TorchDynamo 的日志级别。这些级别包括：

• logging.DEBUG 或 TORCH_LOGS="+dynamo": 打印遇到的每条指令以及下面列出的所有日志级别。

• logging.INFO: 打印每个编译的函数（原始和修改后的字节码）以及捕获的图，此外还包括下面列出的所有日志级别。

• logging.WARNING（默认）：除了下面列出的所有日志级别外，还会打印图表中断。

• logging.ERROR: 仅打印错误。

如果模型非常大，日志可能会变得难以处理。如果在模型的Python代码深处发生错误，执行仅包含错误的帧可能会更有助于调试。有两种工具可以实现这一点：

• 将环境变量 TORCHDYNAMO_DEBUG_FUNCTION 设置为所需函数名称将仅在具有该名称的函数上运行 torchdynamo。

• 启用记录/回放工具（设置 torch._dynamo.config.replay_record_enabled = True）， 当遇到错误时会转储一个执行记录。然后可以回放此记录，仅运行发生错误的帧。

诊断 TorchInductor 错误

如果错误在使用 "eager" 后端时没有发生，那么错误源是后端编译器（错误示例）。 TorchDynamo 有 不同的后端编译器选择，其中 TorchInductor 适合大多数用户的需求。本节以 TorchInductor 为例进行说明，但一些工具也可以与其他后端编译器一起使用。

以下是我们关注的堆栈部分：

使用 TorchInductor 作为后端时，AOTAutograd 用于从 torchdynamo 捕获的前向图中生成反向图。需要注意的是，在此跟踪过程中以及 TorchInductor 将前向图和反向图降低为 GPU 代码或 C++ 代码时，可能会发生错误。一个模型通常由数百或数千个 FX 节点组成，因此确定问题发生的具体节点可能非常困难。幸运的是，有一些工具可以自动将这些输入图缩小到导致问题的节点。第一步是确定错误是在使用 AOTAutograd 跟踪反向图时发生的，还是在 TorchInductor 降低图时发生的。如上文步骤 2 所述，可以使用 "aot_eager" 后端仅在隔离状态下运行 AOTAutograd 而不进行降低。如果在此后端下仍然发生错误，则表明错误发生在 AOTAutograd 跟踪期间。

这是一个例子：

import torch

import torch._dynamo as dynamo

model = torch.nn.Sequential(*[torch.nn.Linear(200, 200) for _ in range(5)])

def test_backend_error():

    y = torch.ones(200, 200)
    x = torch.ones(200, 200)
    z = x + y
    a = torch.ops.aten._foobar(z)  # 虚拟函数，会引发错误
    return model(a)


compiled_test_backend_error = torch.compile(test_backend_error, backend="inductor")
compiled_test_backend_error()

运行此代码应会在其下方显示一个较长的堆栈跟踪，并给出以下错误：

Traceback (最近一次调用最后):
  File "/scratch/mlazos/torchdynamo/torchinductor/graph.py", line 246, in call_function
    return lowerings[target](*args, **kwargs)
  File "/scratch/mlazos/torchdynamo/torchinductor/lowering.py", line 185, in wrapped
    return decomp_fn(*args, **kwargs)
  File "/scratch/mlazos/torchdynamo/torchinductor/lowering.py", line 810, in _foobar
    assert False
AssertionError
...

带有完整堆栈跟踪的错误

如果你将 torch.compile(backend="inductor") 改为 torch.compile(backend="aot_eager")，它将不会报错，因为 问题 出在 TorchInductor 的 lowering 过程中，而不是在 AOTAutograd 中。

缩小TorchInductor错误

从这里开始，让我们运行最小化器以获得一个最小的重现。设置环境变量 TORCHDYNAMO_REPRO_AFTER="aot"（或直接设置 torch._dynamo.config.repro_after="aot"）将生成一个 Python 程序，该程序将 AOTAutograd 生成的图简化为重现错误的最小子图。（请参见下文，我们在其中最小化 TorchDynamo 生成的图）运行带有此环境变量的程序应显示几乎 相同的输出，并额外显示一行指示 minifier_launcher.py 已写入的位置。输出目录可通过将 torch._dynamo.config.base_dir 设置为有效的目录名称来配置。最后一步是运行最小化器并检查它是否成功运行。成功运行看起来像 这样。如果最小化器成功运行，它将生成可运行的 Python 代码，该代码重现了完全相同的错误。对于我们的示例，这是以下代码：

import torch
from torch import tensor, device
import torch.fx as fx
from torch._dynamo.testing import rand_strided
from math import inf
from torch.fx.experimental.proxy_tensor import make_fx

# torch 版本: 1.13.0a0+gitfddfc44
# torch cuda 版本: 11.6
# torch git 版本: fddfc4488afb207971c54ad4bf58130fdc8a4dc5


# CUDA 信息:
# nvcc: NVIDIA (R) Cuda 编译器驱动
# 版权所有 (c) 2005-2022 NVIDIA Corporation
# 构建于 Thu_Feb_10_18:23:41_PST_2022
# Cuda 编译工具, 版本 11.6, V11.6.112
# 构建 cuda_11.6.r11.6/compiler.30978841_0

# GPU 硬件信息:
# NVIDIA A100-SXM4-40GB : 8

from torch.nn import *

class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, add):
        _foobar = torch.ops.aten._foobar.default(add);  add = None
        return (_foobar,)

args = [((200, 200), (200, 1), torch.float32, 'cpu')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
mod = make_fx(Repro())(*args)
from torch._inductor.compile_fx import compile_fx_inner

compiled = compile_fx_inner(mod, args)
compiled(*args)

Repro 模块的 forward 方法包含了导致问题的确切操作。在提交问题时，请包含任何简化后的重现示例以帮助调试。

缩小后端编译器错误

除了TorchInductor之外的后端编译器，查找导致错误的子图的过程与TorchInductor中的错误中的步骤几乎相同，但有一个重要的注意事项。即，最小化器现在将在TorchDynamo跟踪的图上运行，而不是AOTAutograd的输出图。让我们通过一个例子来了解。

import torch

import torch._dynamo as dynamo

model = torch.nn.Sequential(*[torch.nn.Linear(200, 200) for _ in range(5)])
# 玩具编译器，如果图包含relu则失败
def toy_compiler(gm: torch.fx.GraphModule, _):
    for node in gm.graph.nodes:
        if node.target == torch.relu:
            assert False

    return gm


def test_backend_error():
    y = torch.ones(200, 200)
    x = torch.ones(200, 200)
    z = x + y
    a = torch.relu(z)
    return model(a)


compiled_test_backend_error = torch.compile(test_backend_error, backend=toy_compiler)
compiled_test_backend_error()

为了在TorchDynamo跟踪前向图后运行代码，您可以使用TORCHDYNAMO_REPRO_AFTER环境变量。使用TORCHDYNAMO_REPRO_AFTER="dynamo"（或torch._dynamo.config.repro_after="dynamo"）运行此程序应生成此输出以及{torch._dynamo.config.base_dir}/repro.py中的以下代码。

注意

TORCHDYNAMO_REPRO_AFTER 的另一个选项是 "aot"，它将在生成反向图后运行最小化器。

import torch
import torch._dynamo as dynamo
from torch import tensor, device
import torch.fx as fx
from torch._dynamo.testing import rand_strided
from math import inf
from torch._dynamo.debug_utils import run_fwd_maybe_bwd

from torch.nn import *

class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, add):
        relu = torch.relu(add);  add = None
        return (relu,)


mod = Repro().cuda()
opt_mod = torch.compile(mod, backend="None")


args = [((200, 200), (200, 1), torch.float32, 'cpu', False)]
args = [rand_strided(sh, st, dt, dev).requires_grad_(rg) for (sh, st, dt, dev, rg) in args]


with torch.cuda.amp.autocast(enabled=False):
    ref = run_fwd_maybe_bwd(mod, args)
    res = run_fwd_maybe_bwd(opt_mod, args)

压缩器成功地将图表简化为在 toy_compiler 中引发错误的操作。与 TorchInductor 错误 中的过程的另一个不同之处在于，在遇到后端编译器错误后，压缩器会自动运行。成功运行后，压缩器会将 repro.py 写入 torch._dynamo.config.base_dir。

性能分析

访问 TorchDynamo Profiler

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

使用TORCH_COMPILE_DEBUG进行TorchInductor调试

TorchInductor 有一个内置的统计和跟踪功能，用于显示每个编译阶段所花费的时间、输出代码、输出图形可视化和 IR 转储。这是一个调试工具，旨在更容易理解和排查 TorchInductor 的内部问题。

让我们用以下测试程序运行一个示例（repro.py）：

import torch

@torch.compile()
def test_model(x):
    model = torch.nn.Sequential(
        torch.nn.Linear(10, 10),
        torch.nn.LayerNorm(10),
        torch.nn.ReLU(),
    )
    return model(x)


y = test_model(torch.ones(10, 10))

设置环境变量 TORCH_COMPILE_DEBUG=1 将导致创建一个调试跟踪目录，默认情况下，该目录将位于当前目录中，并命名为 torch_compile_debug（这可以在 torchdynamo 配置字段 debug_dir_root 中覆盖，也可以通过环境变量 TORCH_COMPILE_DEBUG_DIR 进行覆盖）。在此目录中，每次运行都将有一个单独的文件夹，文件夹名称由运行的时间戳和进程 ID 组成：

$ env TORCH_COMPILE_DEBUG=1 python repro.py
$ cd torch_compile_debug
$ ls
run_2023_03_01_08_20_52_143510-pid_180167

在运行文件夹中将会有一个torchdynamo目录，其中包含调试日志，以及一个torchinductor文件夹，该文件夹包含每个编译内核的子文件夹，其中包含inductor调试工件。

$ cd
run_2023_03_01_08_20_52_143510-pid_180167
$ ls
torchinductor  torchdynamo

进一步进入 torchinductor 目录，\*.log 文件是编译过程中AOT Autograd阶段的日志，model__0_forward_1.0 包含inductor调试的工件。

$ cd torchinductor
$ ls
aot_model___0_debug.log  模型__0_forward_1.0
$ cd 模型__0_forward_1.0
$ ls
debug.log  fx_graph_readable.py  fx_graph_runnable.py  fx_graph_transformed.py  ir_post_fusion.txt  ir_pre_fusion.txt  output_code.py

以下是内容摘要：

• fx_graph_readable.py 和 fx_graph_runnable.py 是接收到的 fx_graph 的可读和可运行版本。

• fx_graph_transformed.py 是 inductor 运行所有 fx 传递后的 fx 图。

• ir\*.txt 是电感器融合前后的电流。

• output_code.py 是子图的编译后的 Triton 内核。

以下是测试程序的示例调试目录内容：

import torch

@torch.compile()
def test_model(x):
    model = torch.nn.Sequential(
        torch.nn.Linear(10, 10),
        torch.nn.LayerNorm(10),
        torch.nn.ReLU(),
    )
    return model(x)


y = test_model(torch.ones(10, 10))

该调试跟踪中的每个文件都可以通过 torch._inductor.config.trace.* 启用和禁用。由于生成它们的开销较大，默认情况下配置文件和图表都是禁用的。

这种新调试格式中的单个节点看起来像：

buf1: SchedulerNode(ComputedBuffer)
buf1.writes =
    {   MemoryDep(name='buf1', index=0, size=()),
        MemoryDep(name='buf1', index=0, size=(s0,))}
buf1.unmet_dependencies = {MemoryDep(name='buf0', index=c0, size=(s0,))}
buf1.met_dependencies = {MemoryDep(name='primals_2', index=c0, size=(s0,))}
buf1.group.device = cuda:0
buf1.group.iteration = (1, s0)
buf1.sizes = ([], [s0])
class buf1_loop_body:
    var_ranges = {z0: s0}
    index0 = z0
    index1 = 0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('buf0', get_index, False)
        get_index_1 = self.get_index('index0')
        load_1 = ops.load('primals_2', get_index_1, False)
        add = ops.add(load, load_1)
        get_index_2 = self.get_index('index1')
        reduction = ops.reduction('buf1', torch.float32, torch.float32, 'sum', get_index_2, add)
        return reduction

查看 示例调试目录输出 以获取更多示例。

图形断点

给定一个这样的程序：

def some_fun(x):
    ...

compiled_fun = torch.compile(some_fun, ...)
...

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

一些图形中断的原因对于TorchDynamo来说是不可逾越的，并且不容易修复。- 调用除torch之外的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, out_guards, graphs, ops_per_graph, break_reasons, explanation_verbose = (
    dynamo.explain(toy_example, torch.randn(10), torch.randn(10))
)
print(explanation_verbose)
"""
Dynamo 生成了 3 个图，有 2 个图中断和 6 个操作。
 中断原因：
1. call_function BuiltinVariable(print) [ConstantVariable(str)] {}
   文件 "t2.py"，第 16 行，在 toy_example 中
    print("woo")

2. generic_jump
   文件 "t2.py"，第 17 行，在 toy_example 中
    if b.sum() < 0:
 """

输出包括：

• out_guards - 一个包含多个子列表的列表，每个子列表包含必须通过的守卫，以确保追踪的图是有效的。

• graphs - 成功追踪的图模块列表。

• ops_per_graph - 一个嵌套列表，其中每个子列表包含在图中运行的操作。

要在遇到第一个图断裂时抛出错误，请使用 nopython 模式。此模式禁用了 TorchDynamo 的 Python 回退，并且只有在整个程序可以转换为单个图时才会成功。示例用法：

def toy_example(a, b):
   ...

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

过度编译

当 TorchDynamo 编译一个函数（或其中的一部分）时，它会基于对局部变量和全局变量的某些假设，以允许编译器进行优化，并将这些假设表示为在运行时检查特定值的守卫。如果这些守卫中的任何一个失败，Dynamo 将重新编译该函数（或该部分），最多重新编译 torch._dynamo.config.cache_size_limit 次。如果你的程序达到了缓存限制，你首先需要确定哪个守卫失败了，以及你的程序的哪一部分触发了它。

The compile profiler 自动将 TorchDynamo 的缓存限制设置为 1，并在仅观察的“编译器”下运行您的程序，该编译器记录任何保护失败的任何原因。您应该确保运行程序的时间（迭代次数）至少与您遇到问题时运行的时间一样长，并且分析器将在此期间累积统计数据。

如果你的程序表现出有限的动态性，你可能能够调整TorchDynamo缓存限制，以允许每个变化被编译和缓存，但如果缓存限制过高，你可能会发现重新编译的成本超过了任何优化收益。

torch._dynamo.config.cache_size_limit = <你的期望缓存限制>

TorchDynamo 计划支持许多常见的动态张量形状情况，例如变化的批次大小或序列长度。它不计划支持秩动态性。同时，设置特定的缓存限制可以与分桶技术结合使用，以实现一些动态模型的可接受的重编译次数。

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_REPRO_LEVEL=4，精度问题也可以被最小化，它采用类似git bisect的模型，完整的重现可能类似于 TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4，我们需要这个的原因是下游编译器会生成代码，无论是Triton代码还是C++后端，这些下游编译器的数值可能在细微之处有所不同，但会对你的训练稳定性产生重大影响。因此，精度调试器对我们检测代码生成中的错误或后端编译器中的错误非常有用。

如果您希望确保在 torch 和 triton 之间的随机数生成是一致的，那么您可以启用 torch._inductor.config.fallback_random = True