PyTorch 编译栈实战：TorchDynamo 到 TorchInductor 的算子融合
摘要: 从
torch.jit.trace到torch.compile，PyTorch 的编译栈换了三代。本文讲清楚 TorchScript / TorchDynamo / TorchInductor 各自干什么，并用可复现脚本实测torch.compile的加速收益、理解算子融合的原理，以及生产环境该开什么模式。
0. 实验目标
- 梳理 PyTorch 编译栈三代演进：TorchScript → TorchDynamo → TorchInductor
- 理解 TorchDynamo 的图捕获机制（不是 trace 而是 hook 捕获）
- 用
torch.compile实测融合收益，看懂TORCH_LOGS的编译过程 - 掌握生产选型：
mode="reduce-overhead"、fullgraph、动态 shape
1. 三代编译栈：先弄清谁是谁
原理讲解：
- TorchScript（2019）：
torch.jit.script/trace把 Python 模型转成静态图。问题：静态图不支持动态控制流（trace 只录一条路径），Python 特性兼容差——社区口碑一般。 - TorchDynamo（2022）：用 Python frame 的字节码级 hook 自动捕获，只在 Python 层面拦截、把能图化的子图切出来交给后端编译，其余回退 Eager。不需要用户改代码，这就是
torch.compile的第一层。 - TorchInductor（2022 起）：TorchDynamo 的默认后端，把捕获的图翻译成 Triton kernel（GPU）——把 PyTorch 编译栈和 Triton 编译器打通了（这也是本知识库 Triton 编译器 一篇的衔接点）。
text
Python 模型
│ torch.compile
▼
TorchDynamo（字节码级捕获子图）
│
▼
TorchInductor（图 → Triton kernel）
│
▼
CUDA 执行2. 一行加速：torch.compile 开箱实测
python
import torch, time
def model(x):
# 典型 MLP + softmax：由若干小算子组成
for _ in range(4):
x = torch.nn.functional.linear(x, torch.randn(512, 512, device="cuda"))
x = torch.relu(x)
return torch.softmax(x, dim=-1)
x = torch.randn(1024, 512, device="cuda")
# Eager 模式（默认）
for _ in range(5):
model(x)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(200):
model(x)
torch.cuda.synchronize()
t_eager = (time.perf_counter() - t0) / 200
# 编译模式
compiled = torch.compile(model)
for _ in range(5):
compiled(x) # 首次调用触发编译（warmup）
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(200):
compiled(x)
torch.cuda.synchronize()
t_compile = (time.perf_counter() - t0) / 200
print(f"Eager: {t_eager*1000:.2f} ms | Compile: {t_compile*1000:.2f} ms | 加速比 {t_eager/t_compile:.2f}x")实验结果（示例量级 · A100）：
Eager: 0.38 ms | Compile: 0.21 ms | 加速比 1.8x原理讲解（为什么快）：
- Eager 模式下每个算子（linear、relu、softmax）是独立 kernel，一次前向要启动 ~10 个 kernel，每个 kernel 有固定启动开销（~5-10μs）。
- 算子融合：TorchInductor 把相邻算子融合成一个 Triton kernel——linear+relu 融合、softmax 用 reduce 内核合并，kernel 数从 10 降到 3-4，省掉的是 kernel 启动开销和中间张量的读写（显存带宽是算力密集小模型的瓶颈）。
- 层数越多、算子越碎，融合收益越大；单一大 GEMM（如大矩阵乘）本身已接近带宽极限，compile 收益有限——别对"已优化的大算子"期待奇迹。
3. 看穿编译过程：TORCH_LOGS
bash
# 打开编译日志，看 Dynamo 捕获与 Inductor 代码生成
TORCH_LOGS="+dynamo,+inductor" python3 bench.py 2>&1 | head -40
# 关键输出：
# [graph break] —— Python 层无法图化的位置（控制流/外部调用）
# [Generated Triton code] —— Inductor 生成的 kernel 源码工程启示：
- graph break 是编译收益的敌人：图被切得越碎，融合机会越少。常见诱因：Python 控制流（if/for 依赖 tensor 值）、
print、dict操作、外部库调用。 - 用
torch.compile(model, fullgraph=True)强制"要么整图编译、要么报错"——开发期用来暴露 graph break，生产不建议开（动态控制流模型会直接挂）。 - 用
torch._dynamo.explain统计 graph break 位置，逐点修复收益最大：
python
explain = torch._dynamo.explain(fx)(model, x) # 输出 break 原因清单4. 生产选型：模式与动态 shape
python
# 生产推荐组合
compiled = torch.compile(
model,
mode="reduce-overhead", # CUDA graph 减少调度开销，吞吐场景首选
fullgraph=False, # 允许回退，保稳定
)逐段拆解：
mode="reduce-overhead"：额外使用 CUDA Graph 捕获整个计算图一次、重复执行，省掉每次的 kernel 调度开销——小 batch、短序列场景收益明显；大 batch 本身调度占比低，收益趋平。mode="max-autotune"：为每个算子跑 autotune 找最优 kernel 配置，编译时间大幅变长（几分钟），适合"编译一次跑很久"的推理部署；训练迭代场景慎用。- 动态 shape：输入 shape 变化频繁时，Dynamo 会按 shape 重新编译（编译缓存膨胀）。用
torch.compile(..., dynamic=True)或限制 shape 范围减少重编译。 - 训练 vs 推理：训练开 compile 主要省前向+反向的 kernel 启动（吞吐 +10-30%）；推理部署用 compile +
torch.inference_mode()+ CUDA Graph，延迟收益最大。
5. 坑与边界
- 首次调用编译慢：warmup 一次触发编译（秒级到分钟级），生产预热或 CI 预编译缓存（
torch._inductor.cudagraph缓存到磁盘）。 - 动态控制流模型：Transformer 的
if length > N这类 Python 分支会导致 graph break，收益打折——把分支移到模型外，或接受部分编译。 - 自定义算子：第三方 CUDA kernel（如 flash-attn 的某些封装）可能不被 Inductor 识别，graph 会断在这些位置。
- 显存略增：CUDA Graph 捕获需要预留内存，显存紧张场景注意。
- 版本敏感：Dynamo/Inductor 演进快，升级 PyTorch 后 benchmark 需重测（可能有 ±10% 波动）。
6. 工程启示汇总
torch.compile是默认加速手段：小算子密集模型 1.5-2× 起步，零代码改动先试。- 融合是核心机制：kernel 启动开销和中间张量读写是省下来的大头。
- graph break 是敌人：开发期
fullgraph=True暴露、torch._dynamo.explain定位，逐点修。 - 吞吐场景
mode="reduce-overhead"（CUDA Graph），延迟场景配inference_mode()。 - 动态 shape 会膨胀编译缓存：
dynamic=True或约束 shape。 - 编译慢是正常的：预热/缓存进 CI，别在生产首请求时等编译。
- 大算子别期待奇迹：融合收益来自"碎算子"，单一大 GEMM 已接近带宽极限。