PyTorch torch.compile 怎样读取 graph break 诊断信息

当 torch.compile 没有得到预期加速，或者日志里出现多个编译片段时，先不要急着调 Inductor 参数。更直接的做法是打开 graph break 诊断：小例子用 TORCH_LOGS="graph_breaks" 看用户代码位置和原因；需要汇总一次调用中的多个断点时，再用 torch._dynamo.explain。

官方地址：https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html

下面的范围只覆盖“找到断点并判断修复方向”，不把 graph break 诊断扩展成完整性能调优。

先用未编译路径确认是不是编译问题

graph break 日志里的原因不一定代表 TorchDynamo 有缺陷，也可能是函数本身就有参数错误、数据依赖控制流，或调用了暂不支持追踪的 Python 行为。先保留同一组输入，用普通调用做一次基线。

import torch

def forward(x):
    # 让基线和编译路径共用同一段业务逻辑，避免比较两个不同函数
    return torch.sin(x) + 1

compiled_forward = torch.compile(forward, backend="eager")
x = torch.ones(2, 3)

# 先执行未编译版本，确认输入和函数本身能够正常工作
eager_result = forward(x)

# 使用 eager 后端缩短诊断迭代时间，不把问题混到 Inductor 代码生成中
compiled_result = compiled_forward(x)
print(torch.allclose(eager_result, compiled_result))

如果未编译版本也失败，优先修复业务代码；如果只有编译路径暴露断点，再继续读取 Dynamo 的诊断信息。backend="eager" 适合排查追踪边界，因为它减少了后端代码生成带来的噪声。

用 TORCH_LOGS 找到断点的文件和行号

对小程序或已经知道问题函数的场景，环境变量最直观。它会打印 graph break 发生的用户代码位置和原因，通常能直接把排查范围收缩到一个表达式。

# 只打开 graph break 日志，减少其它编译器输出的干扰
TORCH_LOGS="graph_breaks" python diagnose_compile.py

也可以在 Python 入口中打开同一个日志选项：

import torch
import logging

# 只请求 graph break 事件；其它组件保持默认日志级别
torch._logging.set_logs(graph_breaks=True)

@torch.compile(backend="eager")
def forward(x):
    # 这类 Python 输出可能成为追踪边界，正好用来观察诊断信息
    print("shape:", tuple(x.shape))
    return x.sin()

# 触发一次编译，让日志包含真实的用户代码位置
forward(torch.ones(2, 3))

日志里重点看三类信息：用户文件与行号、Reason 后面的原因、以及该位置是否位于模型主要计算路径。不要把每个断点都当成必须消除的错误；预处理或日志函数里的断点，影响可能小于模型 forward 中间的断点。

用 torch._dynamo.explain 汇总一次调用的断点

当一个函数可能出现多个断点时，单看滚动日志不容易比较。torch._dynamo.explain 会运行一次追踪并汇总结果，常用字段包括图数量、断点数量和断点原因。

import torch
import torch._dynamo as dynamo

def forward(x):
    # 用一个会被编译器关注的 Python 行为构造诊断样例
    print("batch:", x.shape[0])
    return x.sin() + 1

# explain 返回的是本次调用的汇总对象，不是模型输出 Tensor
report = dynamo.explain(forward)(torch.ones(2, 3))
print("graph_count:", report.graph_count)
print("graph_break_count:", report.graph_break_count)

# 逐条打印断点原因，便于把修复动作绑定到具体位置
for reason in report.break_reasons:
    print(reason)

看到 graph_break_count 后，不要只记录数字。继续读 break_reasons，确认它指向哪一个函数、哪一行，以及原因是代码错误、数据依赖行为还是当前追踪器暂不支持。不同 PyTorch 版本的文本细节可能变化，判断时应以字段语义和用户代码位置为主。

用 fullgraph=True 让第一个断点更早暴露

默认 fullgraph=False 时，Dynamo 可以编译当前片段、执行暂不支持的代码，再继续追踪，因此一次运行可能留下多个片段。若目标是迅速定位第一个不兼容点，可以临时使用 fullgraph=True。

import torch

@torch.compile(fullgraph=True, backend="eager")
def forward(x):
    # fullgraph=True 会在无法保持单一图时更早报告问题
    if x.sum() > 0:
        return x + 1
    return x - 1

# 用最小输入触发第一个不兼容点，避免直接从完整训练脚本排查
forward(torch.ones(2, 3))

这个模式适合建立最小复现，不代表所有生产函数都必须保持单一图。定位完成后，如果某个函数本来就包含 I/O、日志或其它不适合编译的工作，应把它从编译区域隔离，而不是为了消除日志盲目改写整个模型。

按诊断原因决定修复还是隔离

对高频、低收益的预处理函数，隔离通常比制造大量小图更稳妥；对模型核心计算路径，则应优先按照日志指向改写不支持操作。若日志没有给出可执行建议，却能稳定复现并造成明显性能损失，再准备最小复现和完整诊断信息提交问题。

一份可复用的排查顺序

1. 先用未编译函数运行同一输入，确认业务逻辑没有独立错误。
2. 把后端暂时设为 eager，并开启 TORCH_LOGS="graph_breaks"。
3. 从日志记录用户文件、行号和原因，不只记录断点总数。
4. 多个断点时调用 torch._dynamo.explain，读取 break_reasons。
5. 仍然不清楚边界时，用 fullgraph=True 建立更小的复现。
6. 最后按计算路径收益选择改写、减少日志或隔离函数。

一句话总结：TORCH_LOGS 负责把断点“定位到代码”，explain 负责把一次调用的断点“汇总出来”，而未编译基线与 fullgraph=True 帮助你判断应该修业务代码、改写追踪边界，还是接受一个明确的隔离点。

相关问题

• graph break 一定会让模型变慢吗？不一定，关键看断点是否位于主要计算路径，以及是否造成大量小图和额外调度。
• 为什么只开 graph_breaks 仍看不懂？先用 fullgraph=True 缩小复现，再结合 trace_source 或 trace_bytecode 读取追踪边界。
• 什么时候应该用 disable？当函数包含大量 I/O、日志或暂不支持逻辑，且编译它的收益低于频繁断点开销时，可以隔离该函数。