PyTorch eager 的优点是灵活:每一行 Python 都立刻执行,动态控制流天然可用。 代价是每个 op 都要走 Python 调度、Dispatcher、kernel launch、临时张量分配。 对 GPU 来说,很多小 op 的真正瓶颈不是算力,而是启动开销、内存读写和中间结果搬运。 torch.compile 的核心目标,就是把一段稳定的 PyTorch 计算变成可优化的图, 然后让后端把多个小操作揉成更少、更贴近硬件的 kernel。

它最擅长的三类加速

torch.compile 不是一个单独的优化 pass, 而是一条流水线。粗略看,它把 Python 执行切成可捕获的区域,生成 FX 图, 为这张图生成 guards,再把 forward / backward 交给后端编译。 每个环节都有可能失败、降级或变慢,所以要把链路图刻进脑子里。

组件速查

# 三层调试阶梯:每一层多打开一点编译器能力
compiled = torch.compile(model, backend="eager")       # 只走 Dynamo 捕获,不做后端优化
compiled = torch.compile(model, backend="aot_eager")   # Dynamo + AOTAutograd,仍用 eager 执行图
compiled = torch.compile(model, backend="inductor")   # 默认:完整编译链路

# 用 fullgraph 当作"图完整性测试":有 graph break 就直接报错
compiled = torch.compile(model, fullgraph=True)

graph break 的意思不是"编译失败",而是 Dynamo 捕获到某一行时发现: 这段 Python 行为无法安全放进 FX 图。于是它先把已经捕获的图编译掉, 中间那段 unsupported code 用 eager 执行,然后继续尝试捕获后面的代码。 这保留了 PyTorch 的灵活性,但会打断融合、增加调度边界、降低优化机会。

import torch

@torch.compile
def bad_fn(x):
    y = x.sin() + x.cos()
    if y.sum().item() > 0:       # 数据依赖控制流 + Tensor.item()
        print("positive")       # logging 也会打断图
        return y * 2
    return y - 2

bad_fn(torch.randn(1024, device="cuda"))

# 观察 graph break 位置和原因
# TORCH_LOGS="graph_breaks" python playground.py

常见 graph break 与改法

fullgraph=True 的正确用法

# 推荐结构:把可编译的热路径剥离出来
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.block = torch.compile(Block(), fullgraph=True)

    def forward(self, batch):
        x = preprocess(batch)      # Python / I/O / tokenizer:留在 eager
        x = self.block(x)          # 稳定张量计算:compile
        return postprocess(x)      # 采样 / 日志:留在 eager

编译器要生成快代码,就必须相信一些假设:输入 dtype 是 fp16、shape 是 (32, 128, 768)、stride 连续、module 的某个属性没变。 guards 就是这些假设的 runtime check。后续调用时 guards 通过,直接复用已编译代码; guards 失败,就针对新条件重新编译。重编译次数太多时,Dynamo 会放弃,回到 eager。

静态 shape 与 dynamic=True

诊断命令

# 小样例:看 graph break / guards / recompile 明细
TORCH_LOGS="graph_breaks,guards,recompiles" python day13_compile_bench.py

# 动态 shape 过度特化:看 dynamic shape 相关日志
TORCH_LOGS="dynamic,recompiles" python day13_compile_bench.py --dynamic

# 大模型:生成完整编译报告,再用 tlparse 打开
TORCH_TRACE="/tmp/day13_trace" python train_or_infer.py
pip install tlparse
tlparse /tmp/day13_trace

Dynamo 捕获图只是"拿到了优化材料",真正把图变快的是后端。 PyTorch 默认后端 Inductor 会把高层 PyTorch ops 分解到更底层的 aten / prim 表达, 做 layout、调度、融合和代码生成。GPU 上常见输出是 Triton kernel; CPU 上则生成 C++/OpenMP 代码。理解 Inductor 不必一开始读源码, 先会选模式、看 trace、解释为什么没加速。

compile mode 速查

为什么编译后仍可能不快

Inductor 调试选项

今天的交付物是一份可复现实验记录:同一个小模型分别跑 eager 和 compile, 排除首编译开销,同步 GPU 计时,校验输出正确性,记录 graph break / recompile 日志。 如果你只有 CPU,脚本也能跑,但更推荐在 CUDA 机器上观察 kernel launch 和 fusion 收益。

# day13_compile_bench.py
import argparse
import time
import torch
from torch import nn


class TinyBlock(nn.Module):
    def __init__(self, hidden, expansion):
        super().__init__()
        self.ln = nn.LayerNorm(hidden)
        self.fc1 = nn.Linear(hidden, hidden * expansion)
        self.fc2 = nn.Linear(hidden * expansion, hidden)

    def forward(self, x):
        residual = x
        x = self.ln(x)
        x = self.fc1(x)
        x = torch.nn.functional.gelu(x)
        x = x * 1.41421356 + 0.01
        x = self.fc2(x)
        return x + residual


def sync(device):
    if device.type == "cuda":
        torch.cuda.synchronize()


def bench(fn, x, device, warmup, iters):
    with torch.inference_mode():
        for _ in range(warmup):
            fn(x)
        sync(device)
        start = time.perf_counter()
        for _ in range(iters):
            fn(x)
        sync(device)
    return (time.perf_counter() - start) * 1000 / iters


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", default="default",
                        choices=["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"])
    parser.add_argument("--dynamic", action="store_true")
    parser.add_argument("--fullgraph", action="store_true")
    parser.add_argument("--batch", type=int, default=32)
    parser.add_argument("--seq", type=int, default=128)
    parser.add_argument("--hidden", type=int, default=768)
    parser.add_argument("--iters", type=int, default=100)
    parser.add_argument("--warmup", type=int, default=20)
    args = parser.parse_args()

    torch.manual_seed(0)
    torch.set_float32_matmul_precision("high")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.float16 if device.type == "cuda" else torch.float32

    model = TinyBlock(args.hidden, expansion=4).to(device=device, dtype=dtype).eval()
    x = torch.randn(args.batch, args.seq, args.hidden, device=device, dtype=dtype)

    compiled = torch.compile(
        model,
        mode=args.mode,
        dynamic=args.dynamic or None,
        fullgraph=args.fullgraph,
    )

    with torch.inference_mode():
        eager_out = model(x)
        compiled_out = compiled(x)      # 首次调用会触发 trace + compile
        torch.testing.assert_close(compiled_out, eager_out, rtol=1e-2, atol=1e-2)

    eager_ms = bench(model, x, device, args.warmup, args.iters)
    compile_ms = bench(compiled, x, device, args.warmup, args.iters)
    speedup = eager_ms / compile_ms

    print(f"device={device} dtype={dtype} mode={args.mode} dynamic={args.dynamic}")
    print(f"shape=({args.batch}, {args.seq}, {args.hidden})")
    print(f"eager   : {eager_ms:.3f} ms / iter")
    print(f"compile : {compile_ms:.3f} ms / iter")
    print(f"speedup : {speedup:.2f}x")


if __name__ == "__main__":
    main()

实验命令

# 1) 默认模式,建立 baseline
python day13_compile_bench.py --mode default

# 2) 检查是否有 graph break / recompile
TORCH_LOGS="graph_breaks,recompiles" python day13_compile_bench.py --mode default --fullgraph

# 3) 固定 shape 推理可以试 reduce-overhead
python day13_compile_bench.py --mode reduce-overhead

# 4) 离线压测再试 max-autotune,不要把首次 autotune 算进线上延迟
python day13_compile_bench.py --mode max-autotune

# 5) 变长 seq 对比 dynamic=True 的重编译情况
TORCH_LOGS="recompiles,dynamic" python day13_compile_bench.py --seq 64 --dynamic
TORCH_LOGS="recompiles,dynamic" python day13_compile_bench.py --seq 128 --dynamic

结果记录模板

交付物检查清单

AI Infra 里,编译不是"开关",而是一项容量和延迟工程。 你需要知道它给哪条路径带来收益、给哪条路径带来风险, 以及如何把编译缓存、shape 分桶、灰度验证接进服务系统。

1. 为什么 torch.compile 可以在保持 eager 语义的同时获得优化机会?
2. FX graph 和最终 Triton kernel 是同一个层级的东西吗?中间还经历了哪些阶段?
3. graph break 为什么通常不是 correctness bug,但会影响性能?
4. guards 失败、recompile、fallback 三者的因果关系是什么?
5. 固定 shape 推理服务为什么更适合 reduce-overhead?
6. 如果 benchmark 没加速,你会按什么顺序排查?