PyTorch compile 出现 graph break 怎么定位

使用 torch.compile 后看到 graph break，不代表模型一定编译失败。默认的 fullgraph=False 会先编译已经捕获到的 FX 图，遇到无法追踪的 Python 操作就回到 eager 执行，之后再尝试继续捕获。真正需要定位的是：断点发生在哪里、是否落在模型计算主路径、拆分后的图是否已经小到抵消编译收益。

先让首个 graph break 暴露出来

第一步不是盲目调动态 shape，也不是直接把所有错误都压掉。先在不使用编译的情况下跑一次，确认原始函数本身没有参数、类型或控制流错误；然后把后端换成 eager，减少 Inductor 生成内核带来的等待。

定位阶段可以临时要求整个函数只能变成一张图：

import torch

def score(x):
    y = torch.relu(x)
    # 这个 Python 输出操作不能作为张量图的一部分被捕获
    print("shape:", tuple(y.shape))
    return y.sum()

# eager 只帮助观察 Dynamo 捕获，不把时间花在后端内核生成上
compiled_score = torch.compile(score, backend="eager", fullgraph=True)

try:
    compiled_score(torch.randn(8, 16))
except Exception as exc:
    # 首个断点通常会连同用户代码位置和处理建议一起抛出
    print(type(exc).__name__, exc)

fullgraph=True 的价值是把“可能影响性能的隐式拆分”变成明确错误。它适合缩小范围，不等于生产代码必须消灭所有 Python 逻辑。若日志指向一个只负责记录信息的辅助函数，后面可以单独隔离。

用 explain 和日志缩小范围

知道“有断点”还不够，还要看断点原因和用户代码栈。torch._dynamo.explain 会把一次调用中捕获到的图数量、断点数量、每个断点的原因和涉及的代码位置集中到一个结果里：

import torch
import torch._dynamo as dynamo

def score(x):
    # 用一个可追踪的张量计算作为对照区域
    y = torch.relu(x)
    print("shape:", tuple(y.shape))
    return y.sum()

# explain 只用于定位，不替换正式的编译调用
report = dynamo.explain(score)(torch.randn(8, 16))
print(report)

重点看 Graph Break Count、Break Reasons 以及用户栈，而不是只看异常第一行。比如 builtin: print 说明记录动作本身触发断点；generic_jump 往往提示数据依赖的 Python 分支需要重新设计。示例中的输出是字段形态说明，实际数量会随函数而变。

小函数可以直接打开精确日志：

# 只打印断点位置和原因；需要看重复编译时再追加 recompiles
TORCH_LOGS="graph_breaks,recompiles" python train.py

模型较大时，使用 TORCH_TRACE 生成追踪数据，再交给 tlparse 看 frame、编译区域、断点和重复编译。注意追踪内容可能包含模型代码，不要把敏感项目的完整日志直接公开。

根据断点性质决定重写还是隔离

定位之后把原因分成三类，处理会更稳：

最容易忽略的是循环。一个无关紧要的打印只造成一次切分，通常不值得重构；如果断点处在每个 token 或每个 batch 都会执行的热路径，就要继续追踪。可以先用 profiler 查看编译区域是否被频繁嵌套，再决定是否拆出纯张量函数。

常见问题

graph break 会让 torch.compile 完全失效吗？

不会。默认模式会运行已捕获的图，断点处回到 eager，再尝试继续编译。性能是否受损取决于断点位置、频率和每段图的大小。

为什么 fullgraph=True 一上来就报错？

它要求整个函数进入单张图，任何无法捕获的操作都会立即暴露。这正适合定位，不代表所有业务函数都必须永久使用这个选项。

应该先看 TORCH_LOGS 还是 tlparse？

小样例先看 TORCH_LOGS="graph_breaks"；大型模型先用 tlparse 看全局编译结构，再回到具体函数打开更细日志。