debug-graph-breaks

debug-graph-breaks

描述

通过运行脚本、查看输出、直接应用修复并重新运行验证，诊断并消除 torch.compile 图中断——无论目标是 fullgraph=True，还是仅减少影响最大的中断。

Claude 使用说明

概述

调用此技能时，帮助用户查找、理解并修复其 torch.compile 代码中的图中断。工作流分为三个阶段：检测、诊断、修复。

图中断文档网站 https://meta-pytorch.org/compile-graph-break-site/ 为每种已知的图中断类型提供了详细页面。当用户遇到特定的图中断（例如 GB0059）时，获取相应页面（例如 https://meta-pytorch.org/compile-graph-break-site/gb/gb0059.html），以了解该中断的详细背景、示例和专属修复建议。如果用户提供了图中断文档网站仓库克隆的本地路径，则直接从该目录读取文档文件。

阶段 1：检测——查找图中断

如果用户提供了脚本路径且没有现成的日志，请使用 TORCH_LOGS="graph_breaks" 运行脚本：

TORCH_LOGS="graph_breaks" python your_script.py 2>&1

这会输出每个图中断的以下信息：

• 触发中断的用户代码文件和行号
• 图中断原因（包括 GB 类型、说明和提示）
• 用户代码回溯
• 指向该图中断专属文档的链接

如果输出含糊不清，或无法确定中断的来源，请使用详细模式重新运行：

TORCH_LOGS="graph_breaks" TORCHDYNAMO_VERBOSE=1 python your_script.py 2>&1

详细模式会添加内部 Dynamo 堆栈跟踪和近期字节码指令，有助于厘清常规输出中不明确之处。

如果用户提供了现有的 TORCH_LOGS 输出，直接解析并跳过阶段 1，进入阶段 2。

用户提出要求或提供相关材料时，也可以使用其他检测方法：

选项 B：tlparse 报告

如果用户有 tlparse 报告（由 TORCH_TRACE 生成的结构化 HTML 报告），可以提供：

• tlparse 输出的本地目录路径
• 报告的 URL

tlparse 报告包含：

• index.html：主页面，其中的堆栈树以树状结构显示所有编译，并按状态标注颜色（绿色 = 成功，浅绿色 = 图中断，红色 = 错误）。每次编译都有一个编译 ID，例如 [0/0]（帧 0，首次编译）或 [1/0]（帧 1，首次编译——通常是在图中断后恢复）。
• failures_and_restarts.html：包含所有 RestartAnalysis 事件（图中断）和编译失败的表格，并列出每个事件的完整图中断原因。
• 每次编译对应的目录（例如 -_0_0_0/、-_1_0_1/）：包含生成的构建产物：
  • dynamo_graph_break_reason_N.txt — 完整的图中断详情：用户代码位置、原因、说明、提示、用户堆栈跟踪和内部 Dynamo 回溯。
  • dynamo_output_graph_N.txt — 该次编译生成的 FX 图。
  • compilation_metrics_N.html — 编译时间、图指标（运算、节点、输入）、重启原因、守卫数量和缓存指标。
  • *_ORIGINAL_BYTECODE_N.txt / *_MODIFIED_BYTECODE_N.txt — Dynamo 修改前后的字节码。

阅读 tlparse 报告时：

1. 先查看 failures_and_restarts.html，快速了解所有图中断的摘要。
2. 阅读每个 dynamo_graph_break_reason_N.txt，查看完整详情。
3. 使用 compilation_metrics_N.html 了解各子图对编译时间的影响。
4. 带有 _1 后缀的编译 ID（例如 [0/0_1]）表示发生了重启——_1 表示 Dynamo 在 _0 次尝试中发现图中断后重新开始分析。

选项 C：fullgraph=True（一次修复一个中断）

如果用户希望一次修复一个中断，而不是一次性修复所有中断，请使用 fullgraph=True：

torch.compile(fn, fullgraph=True)(*args)

这会在遇到第一个图中断时引发错误。修复后重新运行，再修复下一个。

阶段 2：诊断 — 了解每个图中断

对于找到的每个图中断，请执行以下操作：

1. 解析图中断原因。 从输出中提取 GB 类型字符串（例如，“Failed to trace builtin operator”“Unsupported Tensor.item() call with capture_scalar_outputs=False”）以及 GB 文档 URL。

2. 获取图中断文档页面。 使用日志输出中的 URL（例如，https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0059.html）获取该图中断类型的详细信息、示例和建议修复方法。如果提供了图中断网站的本地克隆，请改为读取本地目录中的相应文件。

3. 使用日志输出中的提示类别对图中断进行分类：

   • USER_ERROR：用户的代码存在错误，在 eager 模式下也会失败。通过不使用 torch.compile 运行脚本进行验证。
   • FUNDAMENTAL：Dynamo 从根本上无法跟踪此模式（例如，生成器、某些动态控制流）。用户必须重构代码。
   • SUPPORTABLE：Dynamo 理论上可以支持此模式，但目前尚未支持。用户可以提交 PyTorch issue 或寻找变通方法。
   • CAUSED_BY_EARLIER_GRAPH_BREAK：先修复更早发生的图中断；此问题可能会自行消失。
   • DYNAMO_BUG：很可能是 Dynamo 的 bug。建议提交 issue。

4. 确定用户代码中的位置。 使用用户的堆栈跟踪定位导致中断的用户代码的确切行。直接阅读该代码以了解上下文。

5. 阅读用户的脚本。 始终阅读用户脚本中的相关代码，了解周边上下文。将其与图中断文档页面中的示例进行比较，以确定导致中断的确切模式，并判断最有针对性的修复方法。

6. 确定优先级。 并非所有图中断都同等重要。决定先修复哪些问题：

   • 先解决最容易的问题：print() 调用、日志记录、简单的 .item() 调用——这些问题都有直接的修复方法。
   • 优先处理影响最大的情况：热循环或核心前向传递中的图中断，比初始化阶段或不常执行的代码路径中的中断更重要。
   • 最后处理 CAUSED_BY_EARLIER_GRAPH_BREAK：修复更早发生的问题后，这些问题可能会自动消失。
   • 级联中断：循环内的图中断（GB7000）会导致整个帧回退到 eager。修复内部中断即可消除两者。
   • 如果用户不需要 fullgraph=True，请判断哪些中断对性能影响最大，哪些可以暂时保留。

阶段 3：修复 — 消除图中断

针对每个图中断，根据其类型应用修复方法。始终先阅读用户的代码，再直接进行修改。

修复第三方库（例如 HuggingFace transformers）中的图中断

如果图中断源自第三方库（而非用户代码），用户就无法直接修改该库。此时：

1. 检查库的源代码。 如果用户提供了已安装软件包的路径（例如，transformers site-packages 目录），请阅读相关源文件，了解有问题的代码及其调用链。

2. 通过子类进行重写。 创建一个模型子类，并重写导致图中断的方法，使其实现与编译兼容。这样无需修改库，且具有可移植性：

from transformers import BartModel

class CompilableBartModel(BartModel):
    def _override_problematic_method(self, *args, **kwargs):
        # Compile-compatible reimplementation
        ...

model = CompilableBartModel(config)
compiled_model = torch.compile(model, fullgraph=True)

3. 使用 fullgraph=True 进行验证。 应用子类修复后，使用 fullgraph=True 重新运行，以确认所有图中断均已解决。

与 monkey-patching 相比，优先采用此方法，因为它能创建独立且可复用的修复方案。

逃逸方案（编译安全的替代方案）

这些 API 可让代码在编译区域内运行，同时不会导致图中断：

torch._higher_order_ops.print(format_str, *args, **kwargs)

• print() 的编译安全替代方案。使用 Python 格式字符串语法。
• 示例：torch._higher_order_ops.print("Activated shape: {}", h.shape)

torch._dynamo.config.reorderable_logging_functions

• 将日志函数（如 print、logging.info 和自定义日志记录器）添加到此集合中。Dynamo 会对它们重新排序，以避免图中断，同时保留日志调用。
• 示例：torch._dynamo.config.reorderable_logging_functions.add(print)

torch._dynamo.config.ignore_logging_functions

• 将日志函数添加到此列表中，Dynamo 就会在跟踪期间完全跳过它们（这些调用会被丢弃）。
• 示例：torch._dynamo.config.ignore_logging_functions = [print]

@torch._dynamo.decorators.leaf_function

• 装饰器，可让 Dynamo 和 AOT autograd 都将函数视为不透明函数。该函数会在运行时以即时模式运行。需要提供 @fn.register_fake，以便在编译时推断形状。
• 适用场景：函数有副作用、会调用不可编译的库，或者其内部逻辑无法从编译中受益。

@torch._dynamo.decorators.leaf_function
def my_function(x):
    # This runs eagerly at runtime, no graph break
    ...

@my_function.register_fake
def my_function_fake(x):
    # Return a tensor with the correct shape/dtype for compile-time tracing
    return torch.empty_like(x)

• 对于会修改参数的函数：@leaf_function(mutates_args={"buf"})

@torch._dynamo.decorators.nonstrict_trace

• 装饰器，可让 Dynamo 将函数视为不透明函数，但仍由 AOT autograd 跟踪。函数体可以包含会导致图中断的代码，但 AOT autograd 仍会跟踪函数内部，因此其中的算子会由 Inductor 编译。
• 适用场景：函数体会导致图中断，但仍希望 Inductor 优化其中的算子。

@torch._dynamo.decorators.nonstrict_trace
def my_function(x):
    # Dynamo won't trace this, but AOT autograd will
    return x.relu()

torch.compiler.disable()

• 装饰器／上下文管理器，可完全跳过对某个函数的编译。该函数会以即时模式运行，并在调用处导致图中断。
• 适用场景：函数从根本上与编译不兼容，且无需编译。

自定义算子（torch.library）

• 将 Python 函数注册为自定义 PyTorch 算子。Dynamo 会将自定义算子视为图中不透明且类型明确的节点，不会导致图中断，并且 Inductor 可以在运行时调用它们。
• 适用场景：封装形状定义明确的张量输入／输出函数，包括第三方函数或不可跟踪的函数。

import torch.library

@torch.library.custom_op("mylib::my_op", mutates_args=())
def my_op(x: torch.Tensor) -> torch.Tensor:
    # arbitrary non-traceable code here
    return x.numpy().copy()  # example

@my_op.register_fake
def my_op_fake(x):
    return torch.empty_like(x)

• 与 @leaf_function 不同，自定义算子属于公开的 torch.library API，并会作为具名节点显示在 FX 图中，因此更适合在多个编译函数中重复使用的算子。

按图中断类型分类的常见修复方法

print() / logging（GB0059：“Failed to trace builtin operator”）

• 修复选项（按优先顺序排列）：
  1. torch._dynamo.config.reorderable_logging_functions.add(print) — 保留打印，无需更改代码
  2. torch._dynamo.config.ignore_logging_functions = [print] — 编译期间忽略打印
  3. 替换为 torch._higher_order_ops.print("format {}", arg) — 可安全编译的打印 HOP
  4. 如果仅用于调试，则完全移除打印

依赖数据的分支（GB0035、GB0170）

• if tensor.item() > 0: 或 if tensor.sum(): — 控制流取决于张量值。
• 修复方法：简单的 if/else 可使用 torch.cond()，也可以重构代码以避免依赖数据的分支。对于简单情况，可使用 torch.where() 作为无分支替代方案。
• 如果分支判断的值实际上在跟踪时是常量，可使用 torch._check()，或将其设为 Python 常量。

Tensor.item() / tolist()（GB0124、GB0109）

• 这些操作会从张量中提取标量值，从而产生数据依赖。
• 修复选项：
  1. 重构代码，让值保持为张量（例如，使用 torch.norm(h)，而不是调用 h.norm().item() 后再进行标量运算）。
  2. 设置 torch._dynamo.config.capture_scalar_outputs = True，以便在图中捕获 .item()。
  3. 将 .item() 调用移到已编译区域之外。

不支持的函数调用（GB0147）

• Dynamo 无法跟踪的 Python 函数。
• 修复方法：检查是否有对应的 PyTorch 实现。如果调用的是第三方库：
  1. 如果可以进行形状推断，使用带有 register_fake 的 @leaf_function。
  2. 如果函数体包含适合编译的操作，使用 @nonstrict_trace。
  3. 如果函数不需要编译，使用 torch.compiler.disable()。
  4. 将调用移到已编译区域之外。

尝试跟踪生成器（GB0003）

• 生成器无法编译。
• 修复方法：将生成器转换为列表推导式或常规循环，或者在已编译函数之外调用生成器。

不支持的上下文管理器（GB0142）

• 修复方法：将上下文管理器移到已编译区域之外，或检查是否存在兼容 Dynamo 的替代方案（例如，使用 torch.no_grad() 替代 torch.inference_mode()）。

模块级钩子（GB0083、GB0029）

• 模块上的反向钩子需要编译式 autograd。
• 修复方法：启用编译式 autograd（torch._dynamo.config.compiled_autograd=True），或在编译期间移除钩子。

循环中的图中断（GB7000）

• for/while 循环内部发生图中断，会导致整个帧被跳过。
• 修复方法：先修复循环内部导致图中断的根本原因。此类问题始终由另一个图中断导致。

通用策略

• 将不可编译的代码移出 torch.compile： 拆分函数，让可编译的核心代码位于一个函数中，而不可编译的准备和清理代码位于外部。
• 用 torch 等效操作替换 Python 内置函数： 例如，max() -> torch.max()、abs() -> torch.abs()、sorted() -> torch.argsort()。
• 避免在已编译代码中将张量转换为 Python 值： .item()、.tolist()、.numpy()、bool(tensor)、int(tensor)。
• 让张量操作保持为张量操作： 不要提取标量后执行 Python 数学运算，而应直接使用对张量进行操作的 torch 运算（例如，torch.where、torch.clamp、torch.norm）。

迭代工作流

应用修复后，使用 TORCH_LOGS="graph_breaks"（或 fullgraph=True）重新运行脚本以进行验证：

1. 检查图中断是否已解决。
2. 如果出现新的图中断，重复第 2 至第 3 阶段。
3. 目标：要么 fullgraph=True 能无错误运行，要么消除影响最大的图中断。
4. 可选：如果用户是 PyTorch 开发者（可通过 torch 是否在本地构建来判断），且图中断消息不清晰，可建议修改错误消息或创建 issue。

基准测试改进

修复图中断后，帮助用户衡量改进效果。建议采用以下非侵入式方法：

衡量编译时间：

• 使用 torch._dynamo.utils.CompileProfiler 作为后端，查看每个图的编译时间：

with torch._dynamo.utils.CompileProfiler() as prof:
    compiled_fn = torch.compile(fn, backend=prof)
    compiled_fn(*args)
print(prof.report())

• 子图越少 = 编译开销越低。每次图中断都会触发单独的编译。

衡量运行时性能：

• 简单的墙钟时间测量（适用于 CPU 或 CUDA）：

import time

# Warmup
for _ in range(3):
    compiled_fn(*args)
if torch.cuda.is_available():
    torch.cuda.synchronize()

start = time.perf_counter()
for _ in range(100):
    compiled_fn(*args)
if torch.cuda.is_available():
    torch.cuda.synchronize()
elapsed = (time.perf_counter() - start) / 100
print(f"Average iteration: {elapsed*1000:.2f}ms")

比较改进前后：

• 在修复图中断前后运行相同的基准测试。
• 关键指标：子图数量（越少越好）、总编译时间、每次迭代的平均运行时间。
• 可以通过 TORCH_LOGS="graph_breaks" 输出检查子图数量（统计 "Graph break in user code" 行数），也可以查看 tlparse 报告中的堆栈树。

关键参考资料

• 图中断网站：https://meta-pytorch.org/compile-graph-break-site/
• 图中断提示：torch/_dynamo/graph_break_hints.py
• 编译编程模型：https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/compile/programming_model.html
• leaf_function：torch/_dynamo/decorators.py
• nonstrict_trace：torch/_dynamo/decorators.py
• torch.compiler.disable()：torch/_dynamo/decorators.py
• error_on_graph_break()：torch/_dynamo/decorators.py
• torch._higher_order_ops.print：torch/_higher_order_ops/print.py
• CompileProfiler：torch/_dynamo/utils.py