PyTorch torch.compile 怎样定位图断裂原因

如果 torch.compile 开启后没有预期加速，或者日志显示同一个函数被拆成多个 FX graph，先不要直接把问题归因于 Inductor。更稳妥的顺序是：用 eager 基线确认原函数能运行，再用 fullgraph=True 让图断裂尽早暴露，随后通过 TORCH_LOGS=graph_breaks 和 torch._dynamo.explain 锁定触发点。

图断裂不一定是错误：默认 fullgraph=False 时，Dynamo 会编译当前能捕获的 FX graph，执行不可追踪的 Python 代码，再继续编译后面的区域。排查的目标是找出影响模型计算的断裂，而不是机械地消灭每一次断裂。

官方文档：https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/compile/programming_model.dynamo_core_concepts.html

下面按“触发信号—快速判断—处理步骤—回滚路径—告警确认—复盘项”组织一套可执行手册。示例只关注图断裂定位，不讨论具体 GPU 内核的微观调优。

先确认触发信号：慢、拆图或反复编译

常见触发信号有三类：首次编译很慢且后续收益不明显；一个函数中出现多个编译区域；输入形状或参数变化后反复重新编译。图断裂会减少单个图可获得的优化机会，重复编译则会增加编译时间，两者要分开记录。

快速判断：先用 eager 基线排除代码错误

官方排障建议是先确认未编译版本本身可以执行。若 eager 也失败，先修复参数、算子调用或数据准备；否则编译器日志可能把普通 Python 异常包装成图断裂信息，导致排查方向偏移。

import torch

def step(x):
    # 先保留一个足够小的计算函数，便于比较 eager 与 compile。
    y = x.sin() * 2
    return y

x = torch.randn(8, 8)
eager_out = step(x)
print(eager_out.shape)  # eager 能运行，才继续分析编译路径

如果这一步就失败，先不要增加日志选项。把输入、算子和控制流缩到最小后再继续；如果 eager 正常而 torch.compile 变慢或报断裂，再进入下一步。

处理步骤一：用 fullgraph=True 让断点尽快暴露

fullgraph=True 要求整个函数被捕获为一个图；遇到无法捕获的代码时会直接报错，而不是默认拆图后继续。它适合诊断和库兼容性测试，不一定适合直接作为生产配置。

import torch

def forward(x):
    # 这个分支依赖张量值，可能成为图断裂候选。
    total = x.sum()
    if total.item() > 0:
        return x + 1
    return x - 1

debug_forward = torch.compile(forward, fullgraph=True, backend="eager")
try:
    debug_forward(torch.ones(4, 4))
except Exception as exc:
    # 诊断阶段保留异常文本和代码位置，不在这里吞掉错误。
    print(type(exc).__name__, exc)

这里用 backend="eager" 是为了缩短诊断循环：先观察 Dynamo 能否完整捕获，不让后端代码生成失败干扰断点判断。出现异常后，优先看调用栈中第一个属于自己业务文件的行。

处理步骤二：用日志和 explain 读取真实原因

定位到大致函数后，再缩小日志范围。graph_breaks 记录 Dynamo 决定断图的位置和原因，recompiles 则用于观察 guard 失败导致的重新编译；如果只看到后者，不要误把它当成图断裂。

# 只打开图断裂和重新编译信息，减少无关日志干扰。
TORCH_LOGS="graph_breaks,recompiles" python diagnose.py

对小型函数，可以使用 torch._dynamo.explain 汇总断裂原因。不同 PyTorch 版本的返回对象字段可能略有变化，因此脚本应先打印对象摘要，不要把某个内部字段当成长期稳定 API。

import torch

def forward(x):
    # 用数据相关分支制造一个可观察的诊断样本。
    if x.sum().item() > 0:
        return x * 2
    return x / 2

explain = torch._dynamo.explain(forward)
report = explain(torch.ones(3, 3))
# 诊断输出用于发现断裂位置，实际项目应把报告保存到排障记录。
print(report)

判断时重点看三项：断裂发生在哪个业务文件和行号；原因是数据依赖、Python 内置函数、外部 C 扩展还是日志副作用；断裂是否位于模型主计算路径。数据准备、打印和指标上报区的断裂，通常比核心矩阵运算中的断裂更容易接受。

处理步骤三：按原因选择改写、隔离或保留断裂

确定原因后，不要一律加开关。优先把可表达为 PyTorch 操作的逻辑改写为编译器能追踪的形式；数据依赖控制流可考虑合适的高阶算子，.item() 等标量读取则要评估是否真的需要在编译区域内完成。

如果某个函数主要做日志、I/O、预处理或外部库调用，且本身几乎没有张量计算，可以用 torch.compiler.disable 把它明确留在 eager。这样通常比让一个高层函数产生大量细碎图更容易维护。

import torch

@torch.compiler.disable
def report_metrics(value):
    # 日志和指标上报不参与模型计算，显式留在 eager 路径。
    print("metric=", value)

def forward(x):
    # 先完成主要张量计算，再把辅助副作用隔离出去。
    y = x.relu()
    report_metrics(y.mean())
    return y

compiled_forward = torch.compile(forward)
result = compiled_forward(torch.randn(4, 4))
print(result.shape)  # 只用结果形状做最小化的人工检查

若问题来自外部算子或 C 扩展，不能只凭“能跑”就把它强行塞进图中。可以先保持该小段 eager，记录性能收益和兼容性代价，再决定是否寻找 custom operator 或替代实现。

回滚路径：恢复默认拆图并保留最小复现

fullgraph=True 是诊断开关，不是必须长期打开的生产开关。修复后如果仍有一小段无法捕获，但整体吞吐可接受，可以回到默认 fullgraph=False，或只在收益明确的模型函数上启用编译。

回退时保留三份信息：触发输入的形状和 dtype、断裂代码位置、改写前后的日志摘要。这样后续升级 PyTorch 或更换后端时，可以快速判断问题是已消失、位置变化，还是重新出现。

告警确认：不要把每次断裂都当成故障

建议把告警拆成“功能故障”和“性能退化”两级。eager 与 compile 的结果不一致、编译异常未处理、核心计算持续出现大量断裂，应提高等级；仅在一次性预处理、日志或 I/O 区域出现可预期断裂，可记录但不必阻塞发布。

对于输入形状变化造成的重复编译，应另看 recompiles 和 dynamic shapes 选项。图断裂与 guard 失败可能同时存在，但处理动作不同：前者关注可追踪性，后者关注输入约束和动态形状策略。

复盘项：用一张表结束排障

最终判断可以收敛为一句话：先用 eager 证明代码没问题，再用 fullgraph 找位置，用日志和 explain 读原因，最后只优化真正影响主计算的断裂。这样既能避免默认拆图掩盖性能问题，也不会为了追求“单图”而牺牲可维护性。

常见问题

fullgraph=True 一报错是不是 PyTorch 有 bug？ 不一定。它首先说明当前函数不能被完整捕获，可能是数据依赖控制流、Python 副作用或外部扩展；先看业务栈位置和 eager 基线。

图断裂越少越好吗？ 不绝对。核心张量计算中的断裂更值得处理，预处理、日志和 I/O 的断裂可能是合理边界。

TORCH_LOGS=graph_breaks 能解决断裂吗？ 不能，它只提供位置和原因线索；真正修复仍要改写代码、隔离函数或接受 eager 路径。

为什么看到 recompiles 却找不到 graph break？ guard 失败导致的重新编译不等同于图断裂。应分别查看 recompiles 和输入形状、dtype、常量等变化。