""" ============================================================================= CS349E · 第 12 讲 · 01_tracing.py —— 从 Python 到计算图：追踪、特化与数据依赖的分支 ============================================================================= 【这个脚本要回答的问题】 编译器要优化一段 PyTorch / JAX 程序，第一步是"抓图"。最简单的办法是追踪：用假张量调用一遍函数， 把遇到的张量运算记下来。它简单、通用，但 Python 的控制流、副作用、依赖数据值的分支都会出问题。 演示 1：追踪一个小函数：得到的 IR 长什么样；图执行 = 立即执行。 演示 2：追踪"看见"什么、"看不见"什么：Python 循环被展开、Python 的 if 和算术在追踪时就算掉了、 形状被写死、print 只在追踪时发生一次。 演示 3：依赖数据值的分支：报错（jax.jit）、悄悄特化（torch.jit.trace）、加 guard、切断图（graph break，TorchDynamo / Lazy Tensor）。 演示 4：TorchDynamo 看到的是字节码：dis 一下；PyTorch 2 论文里各种抓图方式的稳健性和开销。 【运行】python3 01_tracing.py （< 1 秒） 【注意】演示 4 打印的字节码取决于 Python 版本（作者用 Python 3.13）。 【依赖】numpy；本目录的 graph_lab；../lecture01 的 console。 """ import dis import math import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from graph_lab import F, LazyRunner, TraceError, check_guards, run, trace # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 def gelu_mlp(x, p): """一个小 MLP：x @ W1 → GELU（tanh 近似）→ @ W2。""" h = x @ p["W1"] g = 0.5 * h * (1.0 + F.tanh(math.sqrt(2.0 / math.pi) * (h + 0.044715 * h * h * h))) return g @ p["W2"] def demo1_trace(): print_header("演示 1：追踪一个小函数（x：[4, 8]，W1：[8, 16]，W2：[16, 8]）") rng = np.random.default_rng(0) p = {"W1": rng.standard_normal((8, 16)) / 3, "W2": rng.standard_normal((16, 8)) / 4} g = trace(gelu_mlp, {"x": (4, 8)}, p) print(g) x = rng.standard_normal((4, 8)) eager = gelu_mlp(x, p) # 同一份代码直接用 numpy 数组跑：立即执行（eager） graph = run(g, {"x": x})[0] print(f"\n图执行 vs 立即执行：最大差 {np.abs(graph - eager).max():.1e}；图里 {g.count()} 个运算节点、" f"{sum(1 for n in g.nodes if n.op == 'const')} 个常量。") print_observations([ "追踪 = 用一个\"假张量\"调用函数：x @ W1、0.5 * h、F.tanh(...) 都调用了 Tensor 的运算符重载，每次往图里加一个节点、返回一个新的假张量。" "图就是一个按执行顺序排列的 SSA 列表（每个 %编号 只定义一次）——和 JAX 的 jaxpr、PyTorch 的 FX graph 是同一种东西。", "math.sqrt(2/π) 在 Python 里就算成了一个数，进图的只是一个常量：追踪只记录**张量**运算，其他 Python 代码在追踪时照常执行。", "同一份模型代码既能用 numpy 数组立即执行，也能被追踪——这正是 PyTorch 的设计：eager 模式是默认，torch.compile 在上面抓图。", ]) def demo2_what_is_seen(): print_header("演示 2：追踪\"看见\"什么、\"看不见\"什么") rng = np.random.default_rng(1) p = {"W": rng.standard_normal((8, 8)) / 3} log = [] def f(x, p, n_layers=3, use_residual=True): log.append("调用了一次 Python 函数体") # 副作用：只在追踪时发生 for _ in range(n_layers): # Python 循环：展开 y = F.tanh(x @ p["W"]) x = x + y if use_residual else y # Python 的 if（看的是 Python 值）：只记走到的那一支 T, d = x.shape # x.shape 是 Python 的数：形状被写死 heads = x.reshape(T, 2, d // 2).transpose(1, 0, 2) # 像拆成 2 个头那样变形 return heads.transpose(1, 0, 2).reshape(T, d) * (1.0 / math.sqrt(d)) rows = [] for nl, res in ((1, True), (3, True), (3, False)): g = trace(lambda x, p: f(x, p, nl, res), {"x": (4, 8)}, p) rows.append([nl, "有" if res else "无", g.count(), sum(1 for n in g.nodes if n.op == "matmul"), sum(1 for n in g.nodes if n.op == "add"), sum(1 for n in g.nodes if n.op in ("reshape", "transpose"))]) print_table(["n_layers", "残差", "图里的运算节点", "matmul", "add", "reshape / transpose"], rows) log.clear() g = trace(lambda x, p: f(x, p), {"x": (4, 8)}, p) x5 = rng.standard_normal((5, 8)) try: run(g, {"x": x5}) msg = "没有报错" except ValueError as e: msg = f"报错：{str(e)[:60]}" print(f"\n追踪时 log 里有 {len(log)} 条记录；用追踪出的图跑 3 次，log 不再增加（图里没有 print / append）。") for _ in range(3): run(g, {"x": rng.standard_normal((4, 8))}) print(f"跑完后 log 里还是 {len(log)} 条。用 [4, 8] 追踪的图去跑 [5, 8] 的输入：{msg}") print_observations([ "Python 的 for 循环被**展开**：n_layers = 3 的图里有 3 个 matmul，没有\"循环\"这种节点。层数多的模型因此会得到很大的图" "（JAX 用 lax.scan、PyTorch 用 torch.compile 的层复用来缓解）。", "看 Python 值的 if（use_residual）在追踪时就决定了，图里只有走到的那一支——这叫**特化**（specialization）。" "换一个 use_residual 就得重新追踪一张图。", "副作用（打印、往列表里加东西、修改全局变量）只在追踪时执行一次，之后用图执行时就消失了——TorchDynamo 会把它们识别出来、" "在图外重放（PyTorch 2 论文 §3.7）；torch.jit.trace 直接丢掉。", "形状也被写死了：T、d 在追踪时是 4、8，reshape 的目标形状 (4, 2, 4) 和 1/√8 都成了常量。用 [5, 8] 的输入去跑这张图，" "reshape 就失败了；如果被写死的只是一个常量（比如依赖 T 的缩放），它会悄悄算出错的结果。" "所以编译器要给每张图记下\"它是为哪些形状编译的\"（guard），不符合就重新编译（05 脚本）。", ]) def demo3_data_dependent(): print_header("演示 3：依赖数据值的分支：if x.sum() > 0: ...") def f(x): s = x.sum(axis=-1).sum(axis=-1) if s > 0: # 依赖数据的值！ return x * 2.0 return x * -1.0 rng = np.random.default_rng(2) pos = np.abs(rng.standard_normal((3, 4))) neg = -np.abs(rng.standard_normal((3, 4))) rows = [] try: trace(f, {"x": (3, 4)}, mode="error") rows.append(["报错（像 jax.jit）", "—", "—"]) except TraceError as e: rows.append(["报错（像 jax.jit）", f"TraceError：{str(e)[:30]}…", "要改写成 jnp.where / lax.cond"]) g = trace(f, {"x": (3, 4)}, mode="specialize", example={"x": pos}) err = np.abs(run(g, {"x": neg})[0] - f(neg)).max() rows.append(["悄悄特化（像 torch.jit.trace）", f"用正数追踪，图里只有 x·2（{g.count()} 个节点）", f"负数输入算错，最大误差 {err:.2f}，没有任何警告"]) g = trace(f, {"x": (3, 4)}, mode="guard", example={"x": pos}) ok_pos, ok_neg = check_guards(g, {"x": pos}), check_guards(g, {"x": neg}) rows.append(["记一个 guard", f"guard：s > 0 为真（{len(g.guards)} 个）", f"正数输入 guard {'通过' if ok_pos else '失败'}；负数输入 guard {'通过' if ok_neg else '失败'} → 重新追踪"]) lr = LazyRunner() xt = lr.tensor(neg) s = xt.sum(axis=-1).sum(axis=-1) if s > 0: # 这里 graph break：先执行前半张图，拿到 s 的值 y = xt * 2.0 else: y = xt * -1.0 out = lr.materialize(y) rows.append(["切断图（graph break，像 TorchDynamo / Lazy Tensor）", f"{len(lr.graphs)} 张图，节点数 {lr.graphs}", f"结果正确（最大误差 {np.abs(out - f(neg)).max():.0e}）"]) print_table(["前端的做法", "得到什么", "后果"], rows) print_observations([ "图里的节点只描述\"怎么算\"，不能描述\"看算出来的值决定下一步算什么\"（除非 IR 里有专门的控制流节点，如 lax.cond、torch.cond）。" "追踪时遇到 if 张量:，前端只能三选一。", "torch.jit.trace 的做法最危险：用示例输入走一遍，把走过的路径当成整个函数，换一个输入就悄悄算错——PyTorch 2 论文 §2.1 的例子就是这样。", "graph break：执行到分支时，先把前面记下的那张图编译执行、拿到具体的值，回到 Python 决定走哪条路，再开一张新图。" "结果永远正确，代价是图变小、变多，图之间回到 Python（每次都有开销），跨越分支的融合也做不了了。", "TorchDynamo 对**形状、类型、Python 常量**加 guard，对**张量的值**切断图；PyTorch 2 论文统计：TorchBench 的模型平均每个切成 21.1 张图，70% 的模型没有 graph break。", ]) def demo4_bytecode(): print_header("演示 4：TorchDynamo 看到的东西——字节码；以及各种抓图方式的比较") ins = list(dis.get_instructions(gelu_mlp)) names = [i.opname for i in ins] print(f"gelu_mlp 编译成 {len(ins)} 条字节码（Python {sys.version_info.major}.{sys.version_info.minor}）。前 12 条：") for i in ins[:12]: print(f" {i.offset:4d} {i.opname:<22} {i.argrepr}") print(f"其中二元运算 {sum(1 for n in names if n.startswith('BINARY'))} 条、函数调用 {sum(1 for n in names if 'CALL' in n)} 条。\n") print("PyTorch 2 论文（ASPLOS 2024）表 1、表 2 的数字（A100 / V100，float32）：") print_table(["", "TorchBench", "HuggingFace", "TIMM"], [["模型数", 80, 46, 62], ["TorchDynamo 能抓图", "74（93%）", "46（100%）", "62（100%）"], ["TorchScript 能抓图", "36（45%）", "0（0%）", "61（98%）"], ["每个模型平均几张图", 21.1, 7.7, 1], ["没有 graph break 的模型", "70%", "89%", "100%"]]) print_table(["抓图的额外开销（占 eager 时间）", "推理", "训练"], [["TorchDynamo", "5%", "1%"], ["Lazy Tensors", "38%", "90%"], ["Lazy Tensors + 跨迭代流水", "31%", "86%"]]) print_observations([ "TorchDynamo 不改你的代码，而是在 CPython 执行每个函数之前（PEP 523 的 frame evaluation 钩子）拿到它的字节码，" "**符号地**执行这些字节码：遇到张量运算就记进 FX 图，遇到普通的 Python 运算就照常模拟；遇到处理不了的（数据依赖的分支、" "不认识的库函数）就在那里切断图，生成\"调用图 1 → 执行原来的 Python → 调用图 2\"的新字节码。", "每张图附带一组 guard（论文：约 30 种——张量的形状、dtype、设备、Python 常量、全局状态……），下次调用先检查 guard，" "都通过就直接运行缓存的编译结果，否则重新编译。", "TorchScript 要把 Python 解析成它支持的子集：HuggingFace 的模型返回一个自定义的输出类，它一个都抓不了。" "Lazy Tensors 每次迭代都要在 C++ 里重新记录、哈希整张图，开销大；TorchDynamo 只在第一次分析字节码，之后只检查 guard。", "JAX 的选择是另一种：要求函数是\"纯\"的、用 lax.cond / lax.scan 写控制流，换来的是一张完整的、没有 break 的图（jaxpr）。", ]) if __name__ == "__main__": demo1_trace() demo2_what_is_seen() demo3_data_dependent() demo4_bytecode()