评价此页
★ ★ ★ ★ ★
recipes/debug_mode_tutorial

注意

跳转至页面底部下载完整示例代码。

DebugMode：记录调度操作与数值调试#

作者：Pian Pawakapan, Shangdi Yu

您将学到什么

  • 如何为 Eager 模式和 torch.compile 运行捕获调度操作

  • 如何使用 DebugMode 中的张量哈希和堆栈跟踪来精确定位数值偏差

先决条件

  • PyTorch 2.10 或更高版本

概述#

DebugMode (torch.utils._debug_mode.DebugMode) 是一种 TorchDispatchMode，它会拦截 PyTorch 运行时调用并生成分层操作日志。当你需要了解在 Eager 模式和 torch.compile 下实际运行了什么，或者需要精确定位两次运行之间的数值偏差时，它特别有用。

主要功能

  • 运行时日志记录 – 记录调度操作和 TorchInductor 编译的 Triton 内核。

  • 张量哈希 – 为输入/输出附加确定性哈希值，以便对比运行结果并定位数值偏差。

  • 调度钩子 (Dispatch hooks) – 允许注册自定义钩子来标注调用。

注意

本配方介绍的是一个原型功能。原型功能通常处于早期反馈和测试阶段，可能会有所变动。

快速入门#

下面的代码片段捕获了一个小型 Eager 工作负载并打印出调试字符串。 

from torch._inductor.decomposition import decomps_to_exclude
import torch
from torch.utils._debug_mode import DebugMode

def run_once():
    x = torch.randn(8, 8)
    y = torch.randn(8, 8)
    return torch.mm(torch.relu(x), y)

with DebugMode() as debug_mode:
    out = run_once()

print("DebugMode output:")
print(debug_mode.debug_string())

DebugMode output:
    aten::randn([8, 8], device=cpu, pin_memory=False)  ->  t: f32[8, 8]
    aten::randn([8, 8], device=cpu, pin_memory=False)  ->  t: f32[8, 8]
    aten::relu(t: f32[8, 8])  ->  t: f32[8, 8]
    aten::mm(t: f32[8, 8], t: f32[8, 8])  ->  t: f32[8, 8]

获取更多元数据#

对于大多数调查，你可能需要启用堆栈跟踪、张量 ID 和张量哈希。这些功能提供的元数据可用于将操作关联回模型代码。

DebugMode.log_tensor_hashes 会为每次调用在日志中添加哈希值。hash_tensor 哈希函数使用 torch.hash_tensor，对于元素完全相同的张量，它会返回 0。norm 哈希函数使用 p=1 的范数计算。使用这两个函数（尤其是 norm）时，张量的数值接近度与哈希接近度相关，因此具有相当的可解释性。默认的 hash_fn 是 norm。

with (
    DebugMode(
        # record_stack_trace is only supported for eager in pytorch 2.10
        record_stack_trace=True,
        record_ids=True,
    ) as debug_mode,
    DebugMode.log_tensor_hashes(
        hash_fn=["norm"], # this is the default
        hash_inputs=True,
    ),
):
    result = run_once()

print("DebugMode output with more metadata:")
print(
    debug_mode.debug_string(show_stack_trace=True)
)

DebugMode output with more metadata:
    # File: /var/lib/workspace/recipes_source/debug_mode_tutorial.py:59 in run_once, code: x = torch.randn(8, 8)
    aten::randn([8, 8], device=cpu, pin_memory=False)  ->  t$0: f32[8, 8]  # {'hash': (51.705175606068224,)}

    # File: /var/lib/workspace/recipes_source/debug_mode_tutorial.py:60 in run_once, code: y = torch.randn(8, 8)
    aten::randn([8, 8], device=cpu, pin_memory=False)  ->  t$1: f32[8, 8]  # {'hash': (50.76030981261283,)}

    # File: /var/lib/workspace/recipes_source/debug_mode_tutorial.py:61 in run_once, code: return torch.mm(torch.relu(x), y)
    aten::relu(t$0: f32[8, 8])  ->  t$2: f32[8, 8]  # {'input_hash': (((51.705175606068224,),), {}), 'hash': (27.32746008830145,)}
    aten::mm(t$2: f32[8, 8], t$1: f32[8, 8])  ->  t$3: f32[8, 8]  # {'input_hash': (((27.32746008830145,), (50.76030981261283,)), {}), 'hash': (103.25082559883595,)}

每一行的格式遵循 op(args) -> outputs。启用 record_ids 后，张量将带有 $<id> 后缀，DTensor 将被标记为 dt。

记录 Triton 内核#

尽管 Triton 内核不是通过调度器运行的，但 DebugMode 具有自定义逻辑，可以记录它们的输入和输出。

Inductor 生成的 Triton 内核会显示 [triton] 前缀。前/后哈希标注会报告每次内核调用前后的缓冲区哈希值，这在隔离错误内核时非常有用。

def f(x):
    return torch.mm(torch.relu(x), x.T)

x = torch.randn(3, 3, device="cuda")

with (
    DebugMode(record_output=True) as debug_mode,
    DebugMode.log_tensor_hashes(
        hash_inputs=True,
    )
):
    a = torch.compile(f)(x)

print("Triton in DebugMode logs:")
print(debug_mode.debug_string())

/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_inductor/compile_fx.py:321: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
Triton in DebugMode logs:
    aten::_to_copy(t: f32[3, 3], dtype=torch.float64)  ->  t: f64[3, 3]  # {'input_hash': ((8.278027512133121,), {'dtype': None}), 'hash': 8.278027512133121}
    aten::linalg_vector_norm(t: f64[3, 3], 1)  ->  t: f64[]  # {'input_hash': ((8.278027512133121, None), {}), 'hash': 8.278027512133121}
    aten::_local_scalar_dense(t: f64[])  ->  8.278027512133121  # {'input_hash': ((8.278027512133121,), {}), 'hash': None}
    aten::_to_copy(t: f32[3, 3], dtype=torch.float64)  ->  t: f64[3, 3]  # {'input_hash': ((5.49097703397274,), {'dtype': None}), 'hash': 5.49097703397274}
    aten::linalg_vector_norm(t: f64[3, 3], 1)  ->  t: f64[]  # {'input_hash': ((5.49097703397274, None), {}), 'hash': 5.49097703397274}
    aten::_local_scalar_dense(t: f64[])  ->  5.49097703397274  # {'input_hash': ((5.49097703397274,), {}), 'hash': None}
    [triton] triton_poi_fused_relu_0(in_ptr0=t: f32[3, 3], out_ptr0=t: f32[3, 3], xnumel=9)
    # pre-kernel hashes: {in_ptr0: 8.278027512133121, out_ptr0: 5.49097703397274}
    # post-kernel hashes: {in_ptr0: 8.278027512133121, out_ptr0: 3.2101017609238625}

    aten::_to_copy(t: f32[3, 3], dtype=torch.float64)  ->  t: f64[3, 3]  # {'input_hash': ((8.278027512133121,), {'dtype': None}), 'hash': 8.278027512133121}
    aten::linalg_vector_norm(t: f64[3, 3], 1)  ->  t: f64[]  # {'input_hash': ((8.278027512133121, None), {}), 'hash': 8.278027512133121}
    aten::_local_scalar_dense(t: f64[])  ->  8.278027512133121  # {'input_hash': ((8.278027512133121,), {}), 'hash': None}
    aten::_to_copy(t: f32[3, 3], dtype=torch.float64)  ->  t: f64[3, 3]  # {'input_hash': ((3.2101017609238625,), {'dtype': None}), 'hash': 3.2101017609238625}
    aten::linalg_vector_norm(t: f64[3, 3], 1)  ->  t: f64[]  # {'input_hash': ((3.2101017609238625, None), {}), 'hash': 3.2101017609238625}
    aten::_local_scalar_dense(t: f64[])  ->  3.2101017609238625  # {'input_hash': ((3.2101017609238625,), {}), 'hash': None}
    aten::mm.out(t: f32[3, 3], t: f32[3, 3], out=t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((3.2101017609238625, 8.278027512133121), {'out': 1.4773873090744019}), 'hash': 6.035227698273957}

使用张量哈希进行数值调试#

如果你在不同模式之间发现数值偏差，可以使用 DebugMode 查找偏差源头。在下面的示例中，你可以看到 Eager 模式和编译模式下的所有张量哈希值都是相同的。如果任何哈希值不同，那就是数值偏差出现的地方。

def run_model(model, data, *, compile_with=None):
    if compile_with is not None:
        model = torch.compile(model, backend=compile_with)
    with DebugMode(record_output=True) as dm, DebugMode.log_tensor_hashes(
        hash_inputs=True,
    ):
        dm_out = model(*data)
    return dm, dm_out

class Toy(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x).mm(x.T)

inputs = (torch.randn(4, 4),)
dm_eager, _ = run_model(Toy(), inputs)
dm_compiled, _ = run_model(Toy(), inputs, compile_with="aot_eager")

print("Eager mode:")
print(dm_eager.debug_string())
print("Compiled aot_eager mode:")
print(dm_compiled.debug_string())

Eager mode:
    aten::relu(t: f32[4, 4])  ->  t: f32[4, 4]  # {'input_hash': ((15.85003936290741,), {}), 'hash': 3.1118977069854736}
    aten::permute(t: f32[4, 4], [1, 0])  ->  t: f32[4, 4]  # {'input_hash': ((15.85003936290741, [None, None]), {}), 'hash': 15.85003936290741}
    aten::mm(t: f32[4, 4], t: f32[4, 4])  ->  t: f32[4, 4]  # {'input_hash': ((3.1118977069854736, 15.85003936290741), {}), 'hash': 10.992447525262833}
Compiled aot_eager mode:
    aten::relu(t: f32[4, 4])  ->  t: f32[4, 4]  # {'input_hash': ((15.85003936290741,), {}), 'hash': 3.1118977069854736}
    aten::permute(t: f32[4, 4], [1, 0])  ->  t: f32[4, 4]  # {'input_hash': ((15.85003936290741, [None, None]), {}), 'hash': 15.85003936290741}
    aten::mm(t: f32[4, 4], t: f32[4, 4])  ->  t: f32[4, 4]  # {'input_hash': ((3.1118977069854736, 15.85003936290741), {}), 'hash': 10.992447525262833}

现在让我们看一个张量哈希值不同的示例。我故意写了一个错误的分解，将余弦 (cosine) 分解为了正弦 (sin)。这将导致数值偏差。

from torch._dynamo.backends.common import aot_autograd
from torch._dynamo.backends.debugging import get_nop_func

def wrong_decomp(x):
    return torch.sin(x)

decomp_table = {}
decomp_table[torch.ops.aten.cos.default] = wrong_decomp

backend = aot_autograd(
    fw_compiler=get_nop_func(),
    bw_compiler=get_nop_func(),
    decompositions=decomp_table
)

def f(x):
    y = x.relu()
    z = torch.cos(x)
    return y + z

x = torch.randn(3, 3)
with DebugMode(record_output=True) as dm_eager, DebugMode.log_tensor_hashes(
    hash_inputs=True,
):
    f(x)

with DebugMode(record_output=True) as dm_compiled, DebugMode.log_tensor_hashes(
    hash_inputs=True,
):
    torch.compile(f, backend=backend)(x)

print("Eager:")
print(dm_eager.debug_string(show_stack_trace=True))
print()
print("Compiled with wrong decomposition:")
print(dm_compiled.debug_string())

Eager:
    aten::relu(t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((13.106919825077057,), {}), 'hash': 9.340699851512909}
    aten::cos(t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((13.106919825077057,), {}), 'hash': 4.858168572187424}
    aten::add.Tensor(t: f32[3, 3], t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((9.340699851512909, 4.858168572187424), {}), 'hash': 10.751176416873932}

Compiled with wrong decomposition:
    aten::relu(t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((13.106919825077057,), {}), 'hash': 9.340699851512909}
    aten::sin(t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((13.106919825077057,), {}), 'hash': 7.080562084913254}
    aten::add.Tensor(t: f32[3, 3], t: f32[3, 3])  ->  t: f32[3, 3]  # {'input_hash': ((9.340699851512909, 7.080562084913254), {}), 'hash': 16.421261847019196}

在 Eager 日志中，我们有 aten::cos，但在编译日志中，我们有 aten::sin。此外，Eager 和编译模式下的输出哈希也不同。对比这两个日志将显示第一次数值偏差出现在 aten::cos 调用中。

自定义调度钩子#

钩子允许你使用自定义元数据（例如 GPU 内存使用情况）来标注每次调用。log_hook 返回一个映射，该映射将与调试字符串内联呈现。 

MB = 1024 * 1024.0

def memory_hook(func, types, args, kwargs, result):
    mem = torch.cuda.memory_allocated() / MB if torch.cuda.is_available() else 0.0
    peak = torch.cuda.max_memory_allocated() / MB if torch.cuda.is_available() else 0.0
    torch.cuda.reset_peak_memory_stats() if torch.cuda.is_available() else None
    return {"mem": f"{mem:.3f} MB", "peak": f"{peak:.3f} MB"}

with (
    DebugMode() as dm,
    DebugMode.dispatch_hooks(log_hook=memory_hook),
):
    run_once()

print("DebugMode output with memory usage:")
print(dm.debug_string())

DebugMode output with memory usage:
    aten::randn([8, 8], device=cpu, pin_memory=False)  ->  t: f32[8, 8]  # {'mem': '8.125 MB', 'peak': '14.128 MB'}
    aten::randn([8, 8], device=cpu, pin_memory=False)  ->  t: f32[8, 8]  # {'mem': '8.125 MB', 'peak': '8.125 MB'}
    aten::relu(t: f32[8, 8])  ->  t: f32[8, 8]  # {'mem': '8.125 MB', 'peak': '8.125 MB'}
    aten::mm(t: f32[8, 8], t: f32[8, 8])  ->  t: f32[8, 8]  # {'mem': '8.125 MB', 'peak': '8.125 MB'}

模块边界#

record_nn_module=True 会插入 [nn.Mod] 标记，显示每个操作集由哪个模块执行。截至 PyTorch 2.10，它仅在 Eager 模式下有效，但对编译模式的支持正在开发中。

class Foo(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.l1 = torch.nn.Linear(4, 4)
            self.l2 = torch.nn.Linear(4, 4)

        def forward(self, x):
            return self.l2(self.l1(x))

class Bar(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.abc = Foo()
        self.xyz = torch.nn.Linear(4, 4)

    def forward(self, x):
        return self.xyz(self.abc(x))

mod = Bar()
inp = torch.randn(4, 4)
with DebugMode(record_nn_module=True, record_output=False) as debug_mode:
    _ = mod(inp)

print("DebugMode output with stack traces and module boundaries:")
print(debug_mode.debug_string(show_stack_trace=True))

DebugMode output with stack traces and module boundaries:
    [nn.Mod] Bar
      [nn.Mod] Bar.abc
        [nn.Mod] Bar.abc.l1
          aten::t(t: f32[4, 4])
          aten::addmm(t: f32[4], t: f32[4, 4], t: f32[4, 4])
        [nn.Mod] Bar.abc.l2
          aten::t(t: f32[4, 4])
          aten::addmm(t: f32[4], t: f32[4, 4], t: f32[4, 4])
      [nn.Mod] Bar.xyz
        aten::t(t: f32[4, 4])
        aten::addmm(t: f32[4], t: f32[4, 4], t: f32[4, 4])

结论#

在本教程中，我们了解了 DebugMode 如何提供一个轻量级的、仅限运行时的视图，让你查看 PyTorch 实际执行了什么，无论你是在运行 Eager 代码还是编译后的图。通过叠加张量哈希、Triton 日志记录和自定义调度钩子，你可以快速追踪数值差异。这在调试运行之间的逐位等价性时特别有用。

脚本总运行时间：（0 分 2.638 秒）

在此页面上
PyTorch 库