PyTorch 2.0 故障排查

作者: Michael Lazos

我们正在积极开发调试工具和性能分析器,并不断优化错误和警告信息。以下是各个工具及其常见用途的列表。如需进一步的帮助,请参考诊断运行时错误

标题

工具

目的

使用方法

信息logging

查看编译的简化步骤

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 的 minifier

如果错误在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)。堆栈中的任何组件都可能发生错误,并且会提供完整的堆栈跟踪。

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

步骤

组件

1

TorchDynamo

2

编译器后端

3

Torch Inductor

如果信息日志不够详细,可以使用可用的后端选项。这些选项包括:

  • "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的日志级别。这些级别的具体选项包括:

  • 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 错误中的过程相同,但有一个重要区别:minifier 现在会在 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 profiler

TorchDynamo 内置了一个统计功能,用于收集和显示每个编译阶段所花费的时间。执行 TorchDynamo 后,可以通过调用 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  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 是在运行所有 fx passes 之后,由 inductor 生成的 fx 图。

  • ir\*.txt 包含电感器 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引入必要的保护措施(参见使Dynamo安全:保护措施),以确保编译后的程序可以安全地重用。如果图中断导致生成的片段很小,这会阻碍性能。为了最大化性能,尽量减少图中断的数量是很重要的。

确定图中断的原因

要识别程序中所有图中断及其原因,可以使用 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_verbose)
"""
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:
  ...
"""

输出包括以下内容:

  • out_guards - 一个列表的列表,每个子列表包含确保跟踪图有效的必需守卫。

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

  • ops_per_graph - 一个列表的集合,每个子列表包含了图中运行的操作。

要在遇到第一个图中断时抛出错误,请使用fullgraph模式。此模式禁用了TorchDynamo的Python回退,并且仅在程序可以完全转换为单个图的情况下才成功。示例用法:

def toy_example(a, b):
   ...

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

过度重新编译

当 TorchDynamo 编译一个函数(或其部分)时,会对局部变量和全局变量做出某些假设以允许编译器优化,并将这些假设表达为运行时检查特定值的保护条件。如果任何保护条件失败,Dynamo 将重新编译该函数(或其部分),最多达到 torch._dynamo.config.cache_size_limit 次。如果你的程序达到了缓存限制,你需要首先确定哪个保护条件失败以及是哪一部分代码触发了它。

如果你的程序具有有限的动态性,可以调整TorchDynamo的缓存限制,以便每个变化都能被编译和缓存。但是,如果缓存限制设置得太高,你会发现重新编译的成本超过了优化所带来的收益。

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

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

准确性的调试

通过设置环境变量 TORCHDYNAMO_REPRO_LEVEL=4,可以减少准确性问题。它采用类似 git bisect 的模型,并且完全重现可能类似于 TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4。我们需要这个的原因是下游编译器会生成代码(无论是 Triton 代码还是 C++ 后端代码),这些编译器在细微之处的数值可能有所不同,但会对训练稳定性产生重大影响。因此,准确性调试器对于我们检测代码生成或后端编译器中的错误非常有用。

如果你希望确保 torch 和 triton 中的随机数生成一致,可以将 torch._inductor.config.fallback_random = True 开启。

扩展调试

可以使用以下实验性标志启用扩展调试。

TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED - 当守卫的字符串表示与该标志值匹配时,提供详细的调试信息。例如,设置为“Ne(s0, 10)”以在每次发出守卫时生成完整的Python和C++调用栈。TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL - 当特定符号被分配时,提供详细的调试信息。例如,将其设置为“u2”,以便每当创建此符号时都生成完整的Python和C++调用栈。TORCHDYNAMO_EXTENDED_DEBUG_CPP - 为所有扩展调试设置以及错误提供详细的调试信息(包括C++ 调用栈)。例如,将其设置为“1”。由于C++ 调用栈很慢且会产生大量日志输出,默认情况下不会包含在扩展调试中。

冷启动.timing和缓存损坏调试

为了使句子更加自然流畅,可以调整为:

冷启动时序分析及缓存损坏调试

不过根据要求只返回答案本身,所以直接修改原文中的“Timing”为更符合中文表达的词汇:“时序分析”,最终结果如下:

冷启动时序分析和缓存损坏调试

为了测量冷启动编译时间或调试缓存损坏问题,可以传递环境变量 TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 或设置 torch._inductor.config.force_disable_caches = True。这将覆盖所有其他缓存配置选项,并禁用所有编译时间的缓存。

本页目录