快捷方式

PyTorch 2.0 故障排除

作者: Michael Lazos

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

标题

工具

用途

用法

信息日志记录

查看编译的摘要步骤

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 准确性最小化器

找到在您怀疑问题出在后端(例如,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=<name>

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
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._logging.set_logs(dynamo = logging.INFO)TORCH_LOGS="dynamo" 设置 TorchDynamo 的 log_level。这些级别包括

  • logging.DEBUGTORCH_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)  # dummy function which errors
    return model(a)


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

运行此代码应该会给出此错误,并在其下方显示更长的堆栈跟踪

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)

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)])
# 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 中的以下代码。

注意

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 分析器

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 以及 env var 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  model__0_forward_1.0
$ cd model__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.pyfx_graph_runnable.py 是 inductor 收到的 fx_graph 的可读和可运行版本。

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

  • ir\*.txt 是 inductor 融合前后的 ir。

  • 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 produced 3 graphs, with 2 graph breaks 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 次。如果您的程序达到缓存限制,您首先需要确定哪个守卫失败以及程序的哪一部分触发了它。

编译分析器 自动执行将 TorchDynamo 的缓存限制设置为 1 并运行程序的过程,该程序在仅观察的“编译器”下运行,该编译器记录任何守卫失败的原因。您应该确保运行您的程序至少与遇到问题时运行的时间(迭代次数)一样长,分析器将在此期间累积统计信息。

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

torch._dynamo.config.cache_size_limit = <your desired cache 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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源