TorchDynamo 故障排除

作者: Michael Lazos

TorchDynamo 仍在积极开发中，许多导致图中断和过度重新编译的原因将在即将支持追踪动态张量形状、更谨慎地选择保护措施以及更好的启发式调整后得到修复。

与此同时，你可能需要诊断某个具体问题，并判断是通过修改模型来轻松解决，还是需要提交问题以获取支持。

此外，我们正在积极开发调试工具、性能分析器，并改进我们的错误/警告。如果您对这些基础设施有任何问题或改进建议，请给我们反馈。下表列出了可用的工具及其典型用法。有关更多帮助，请参见 诊断运行时错误。

标题

工具	目的	使用
信息日志记录	查看编译的简要步骤	torch._dynamo.config.log_level = logging.INFO
调试日志	查看编译的详细步骤（打印每个追踪的指令）	torch._dynamo.config.log_level = logging.DEBUG 和 torch._dynamo.config.verbose = True
适用于任何后端的压缩器	查找能够重现任何后端错误的最小子图	设置环境变量 TORCHDYNAMO_REPRO_AFTER="dynamo"
压缩器用于 TorchInductor	如果错误已知在 AOTAutograd` 之后发生 查找能够重现 TorchInductor 降级期间错误的最小子图	设置环境变量 TORCHDYNAMO_REPRO_AFTER="aot"
Dynamo 精度最小化器	查找在怀疑问题出在 AOTAutograd 时，能够重现急切模式模型与优化模型之间精度问题的最小子图	TORCHDYNAMO_REPRO_AFTER="dynamo" TORCHDYNAMO_REPRO_LEVEL=4
Inductor accuracy minifier	查找在怀疑问题出在后端（例如 inductor）时，能够重现急切模式模型与优化模型之间精度问题的最小子图。 如果此方法不起作用，请改用 Dynamo 精度最小化工具。	TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4
torch._dynamo.explain	查找图断点并显示其原因	torch._dynamo.explain(fn, *inputs)
Record/Replay	记录并重放用于在图捕获期间重现错误的帧	torch._dynamo.config.replay_record_enabled = True
TorchDynamo 函数名称过滤	仅编译具有给定名称的函数，以在调试问题时减少噪声	设置环境变量 TORCHDYNAMO_DEBUG_FUNCTION=<name>
TorchInductor 调试日志记录	打印通用 TorchInductor 调试信息和生成的 Triton/C++ 代码	torch._inductor.config.debug = True
TorchInductor 跟踪	显示每个 TorchInductor 阶段所用时间 + 输出代码和图可视化	设置环境变量 TORCH_COMPILE_DEBUG=1 或 torch._inductor.config.trace.enabled = True

诊断运行时错误

以下是 TorchDynamo 编译器栈。

从高层次来看，TorchDynamo 堆栈由从 Python 代码中捕获图（TorchDynamo）和一个后端编译器组成。在此示例中，后端编译器包括反向图追踪（AOTAutograd）和图降级（TorchInductor）*。堆栈中的任何组件都可能出现错误，并会提供完整的堆栈跟踪。

您可以使用信息日志记录 (torch._dynamo.config.log_level = logging.INFO) 并查找 Step #: ... 的输出，以确定错误发生在哪个组件中。日志会在每一步的开始和结束时生成，因此错误对应的步骤应为最近一次已记录但尚未结束的步骤。这些步骤对应于堆栈中的以下部分（根据上图）：

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

AOTAutograd 的开始和结束目前没有记录，但我们计划很快添加这一功能。

如果信息日志记录不足，还有一些后端选项可以帮助你确定是哪个组件导致了错误，特别是当你无法理解生成的错误信息时。这些选项如下：

• "eager": 仅运行 torchdynamo 正向图捕获，然后 使用 PyTorch 运行捕获的图。这可以表明 TorchDynamo 是否引发了错误。

• "aot_eager": 运行 torchdynamo 来捕获一个前向图，然后使用 AOTAutograd 来追踪反向图，而无需任何额外的后端编译步骤。PyTorch 的 eager 模式将用于运行前向和反向图。这有助于将问题缩小到 AOTAutograd。

缩小问题的一般步骤如下：

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

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

这些情况将在以下各节中进行分析。

注意

The TorchInductor backend consists of both AOTAutograd tracing and the TorchInductor compiler itself. We will disambiguate by referring to TorchInductor as the backend, and TorchInductor lowering as the phase which lowers the graph traced by 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
due to:
Traceback (most recent call last):
  File "/scratch/mlazos/torchdynamo/torchdynamo/symbolic_convert.py", line 837, in BUILD_MAP
    assert isinstance(k, ConstantVariable) or (
AssertionError

from user code:
   File "/scratch/mlazos/torchdynamo/../test/errors.py", line 34, in test_assertion_error
    z = {y: 5}

Set torch._dynamo.config.verbose=True for more information
==========

如消息所示，您可以将 torch._dynamo.config.verbose=True 设置为获取完整的堆栈跟踪，包括 TorchDynamo 中的错误和用户代码中的错误。除了此标志外，您还可以通过 torch._dynamo.config.log_level 设置 torchdynamo 的 log_level。可用级别如下：

• logging.DEBUG: 打印遇到的每条指令，以及所有以下日志级别。

• logging.INFO: 打印每个被编译的函数（原始和修改后的字节码） 以及捕获的图，同时包含所有以下日志级别。

• logging.WARNING (默认): 除所有低于此日志级别的信息外，还会打印图中断信息。

• logging.ERROR: 仅打印错误。

如果模型足够大，日志可能会变得令人难以应对。如果在模型的 Python 代码深处发生错误，仅执行发生错误的帧可能会有助于更轻松地进行调试。有两种工具可用于实现这一点：

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

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

TorchInductor 错误

如果错误没有在"eager"后端发生，则后端编译器是错误的来源（示例错误）。 对于TorchDynamo，有不同的选择 适用于大多数用户的后端编译器，如TorchInductor或nvfuser。本节以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)  # dummy function which errors
    return model(a)


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

运行此代码应该会给你这个错误，并在下面显示更长的堆栈跟踪 it:

Traceback (most recent call last):
  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 降低过程中，而不是在 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 version: 1.13.0a0+gitfddfc44
# torch cuda version: 11.6
# torch git version: fddfc4488afb207971c54ad4bf58130fdc8a4dc5


# CUDA Info:
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2022 NVIDIA Corporation
# Built on Thu_Feb_10_18:23:41_PST_2022
# Cuda compilation tools, release 11.6, V11.6.112
# Build cuda_11.6.r11.6/compiler.30978841_0

# GPU Hardware Info:
# 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)

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

压缩后端编译器错误

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

import torch

import torch._dynamo as dynamo

model = torch.nn.Sequential(*[torch.nn.Linear(200, 200) for _ in range(5)])
# toy compiler which fails if graph contains 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中。

注意

The other option for TORCHDYNAMO_REPRO_AFTER are "aot", which will run the minifier after the backward graph has been generated.

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)

The minifier successfully reduced the graph to the op that raises the error in toy_compiler. The other difference from the procedure in TorhInductor Errors is that the minifier is automatically run after encountering a backend compiler error. After a successful run, the minifier writes repro.py to torch._dynamo.config.base_dir.

性能分析

访问 TorchDynamo 分析器

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

TorchInductor 调试追踪

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

设置环境变量 TORCH_COMPILE_DEBUG=1 将导致创建并打印一个 调试跟踪目录:

$ env TORCH_COMPILE_DEBUG=1 python repro.py
torch._inductor.debug: [WARNING] model_forward_0 debug trace: /tmp/torchinductor_jansel/rh/crhwqgmbqtchqt3v3wdeeszjb352m4vbjbvdovaaeqpzi7tdjxqr.debug

这里是一个示例调试目录输出 用于测试程序：

torch.nn.Sequential(
        torch.nn.Linear(10, 10),
        torch.nn.LayerNorm(10),
        torch.nn.ReLU(),
    )

每个调试跟踪中的文件都可以通过 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 = dynamo.explain(toy_example, torch.randn(10), torch.randn(10))
print(explanation)
"""
Dynamo produced 3 graphs, with 2 graph break and 6 ops.
 Break reasons:
1. call_function BuiltinVariable(print) [ConstantVariable(str)] {}
   File "t2.py", line 16, in toy_example
    print("woo")

2. generic_jump
   File "t2.py", line 17, in 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 automates the process of setting TorchDynamo’s cache limit to 1 and running your program under an observation-only ‘compiler’ that records the causes of any guard failures. You should be sure to run your program for at least as long (as many iterations) as you were running when you ran into trouble, and the profiler will accumulate statistics over this duration.

如果你的程序表现出有限的动态性，你可能可以通过调整 TorchDynamo 的缓存限制，使每种变化都能被编译和缓存。但如果缓存限制过高，你可能会发现重新编译的代价超过了任何优化带来的好处。

torch._dynamo.config.cache_size_limit = <your desired cache limit>

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

from torch._dynamo.utils import CompileProfiler

prof = CompileProfiler()

def my_model():
    ...

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++后端，这些下游编译器的数值计算可能会有细微的不同，但会对 你的训练稳定性产生显著影响。因此，精度调试器对我们检测代码生成中的错误或后端编译器的问题非常有用。

提交问题

如果您在使用TorchDynamo时遇到问题，请在GitHub上提交问题。

在提交问题之前，请阅读 README、 TROUBLESHOOTING，并搜索类似的问题。

在提交问题时，请通过运行以下命令包含有关你的操作系统、Python、PyTorch、CUDA 和 Triton 版本的信息：

python tools/verify_install.py

• 如果可能，请提供一个可以重现问题的最小脚本，可以通过运行 Minifier 生成

• 对错误的描述

• 预期行为

• A log (set torch._dynamo.config.log_file to a valid file name to dump the logs to a file and torch._dynamo.config.log_level = logging.DEBUG and torch._dynamo.config.verbose = True)