torch.compile：编译后端Backends

可视化演示

展示torch.compile中Dynamo捕获FX图后Inductor两阶段lower到IR再codegen的编译流程、inductor/eager/aot_eager/cudagraphs四类后端的优化与加速对比、以及按场景的后端选型策略。

后端概览：厨房里的四种流派

小美的厨房新装了“智能大厨系统”（torch.compile），但绑定哪位师傅，得看接什么单子。所谓“编译后端”，就是这个系统把模型翻译成真正可执行程序时使用的底层引擎——你可以把它想象成大厨的“做菜流派”。PyTorch 2.0 通过一行代码就能切换师傅：

参数 backend 就是点名用哪位师傅。四个常任选项分别是：inductor（全能主厨）、eager（老实快炒）、aot_eager（计划派）、cudagraphs（录像大师）。他们做的是同一道菜——同一个模型，区别在于备菜策略、炒菜节奏和额外服务费（编译时间）。小美刚开始总想选“最贵最全能”的，但导师提醒她：选后端要看总账，就像请大厨不能只看手艺，还得等得起他研究菜谱的时间。

$$\displaystyle T_{\text{total}} = T_{\text{compile}} + N \times T_{\text{step}}$$

1. 总耗时定义：$T_{\text{total}}$ 表示跑完整个训练或大规模推理任务的总时间，是选后端的最终裁判。

$$T_{\text{total}}$$

2. 一次性编译开销：$T_{\text{compile}}$ 是模型首次编译的固定成本，类似大厨研究菜谱的“备课时间”，所有后端启动时支付一次。

$$T_{\text{compile}}$$

3. 迭代次数：$N$ 是训练步数或推理请求次数。$N$ 越大，运行阶段的权重越高，编译投资越容易回本。

$$N$$

4. 每步耗时：$T_{\text{step}}$ 是平均每一步（一次前向或反向）的执行时间，不同后端在此天差地别。

$$T_{\text{step}}$$

小美把这条公式贴在厨房墙上，每次换后端先算一笔账：10 次迭代的小任务，别麻烦全能主厨；一万次的大单，一定要把最强的师傅请出来。

torch.compile(model, backend="inductor")

Inductor后端：全能主厨的高效流水线

小美的餐厅接到一个大单——给一万桌客人做同款宴席。这时全能主厨 inductor 登场了。他不满足于照着原菜谱一铲一铲炒，而是把整本菜谱吃透后重新编排：能一起下锅的食材合并处理，能提前备好的料批量切好，甚至自带一口“特制炒锅”（Triton 编译器）为每道菜定制最高效的颠勺动作。这对应到技术上，就是 inductor 会为 GPU 生成专门的优化内核，并把多个算子自动融合。你可以把它理解为：把“加水→搅拌→加热”三个动作合并成一个“一体机”操作，中间不再把食材端进端出，省掉了大量读写内存的时间。

$$\displaystyle S = \frac{\sum_{i=1}^{n} t_i}{t_{\text{fused}} + t_{\text{launch}}}$$

1. 拆分执行耗时：$\sum_{i=1}^{n} t_i$ 是原始模型中 n 个算子各自执行的时间累加，每个算子都要单独从显存读数据、算完再写回去。

$$\sum_{i=1}^{n} t_i$$

2. 融合后计算：$t_{\text{fused}}$ 是融合内核实际完成全部计算的时间，因为数据留在寄存器/缓存里，省去大量访存，通常远小于拆分总和。

$$t_{\text{fused}}$$

3. 单次启动成本：$t_{\text{launch}}$ 是 CPU 把内核“发号施令”给 GPU 的一次性调度开销，从 n 次降到 1 次。

$$t_{\text{launch}}$$

inductor 之所以是 PyTorch 2.0 的默认后端，正是因为它在这条公式上表现最出色。但它的代价是“备课”时间最长——编译经常要花几十秒甚至几分钟。所以小美只在大批量、长周期任务上请他出山，一次研究，万次受益。

Eager后端：随炒随吃的老实厨子

从前没有智能大厨的时候，小美自己是厨师：接到一单，立刻开火，炒完一道菜端出去，再回来做下一道。eager 后端就是这位“老实厨子”的数字化版本——PyTorch 早期默认的即时执行模式。每个算子（比如矩阵乘法、激活函数）被 Python 遇到时，立刻在 GPU 上执行，绝不提前规划。好处显而易见：零编译时间，拿来就用；出了 bug 可以像跟厨子对话一样，逐步打断、审查每一道菜（算子）的输入输出，调试体验极佳。

$$\displaystyle T_{\text{step}}^{\text{eager}} = \sum_{i=1}^{m} \left( t_i^{\text{launch}} + t_i^{\text{compute}} \right)$$

1. 算子数量：$m$ 是模型中算子的总数，像一份菜谱里的小步骤数，步骤越多，逐项执行的浪费越大。

$$\sum_{i=1}^{m}$$

2. 单算子启动代价：$t_i^{\text{launch}}$ 是第 i 个算子被 CPU 调度到 GPU 的开销。eager 模式每步都要“喊一嗓子”，喊多了自然累。

$$t_i^{\text{launch}}$$

3. 单算子计算代价：$t_i^{\text{compute}}$ 是第 i 个算子真正算矩阵的时间。算得快，但调度拖后腿。

$$t_i^{\text{compute}}$$

正因如此，eager 的每一步都比“融合派”要慢。但当模型还在频繁修改、或需要精确定位某个数值异常时，这位老实厨子反而是最可靠的伙伴——他从不藏私，每一铲子你都看得清清楚楚。

AOT Eager后端：提前备菜的计划派厨师

aot_eager 这位厨师很特别：他会在开工前把整本菜谱从头到尾翻一遍，把“先放油还是先爆香”这类流程问题全部定下来，写成一页流程单（计算图），然后才开始一铲一铲现炒。他并不像 inductor 那样合并食材、定制锅具，只是省去了炒菜过程中反复翻菜谱的犹豫时间。技术上说，aot_eager 利用 AOTAutograd 在编译期把前向和反向的计算图完整展开并缓存起来，运行时不需要像 eager 那样动态追踪、临时构建执行图，省掉了 tracing 的 CPU 开销。

$$\displaystyle T_{\text{step}}^{\text{aot}} = T_{\text{setup}} + \sum_{i=1}^{m} \left( t_i^{\text{launch}} + t_i^{\text{compute}} \right)$$

1. 预建图开销：$T_{\text{setup}}$ 是运行期取出预缓存计算图的固定成本，因为它已在编译期展开，这里几乎可以忽略。

$$T_{\text{setup}}$$

2. 算子启动与计算：$\sum_{i=1}^{m} \left( t_i^{\text{launch}} + t_i^{\text{compute}} \right)$ 与 eager 相同，仍然逐算子执行，不做融合——这就是它与 inductor 的本质区别。

$$\sum_{i=1}^{m} \left( t_i^{\text{launch}} + t_i^{\text{compute}} \right)$$

计划派师傅的价值在于：当模型计算图结构复杂、动态建图的 CPU 开销明显时，他能帮你省下这部分冤枉时间，同时保留逐算子执行的透明性。小美常把它当作 eager 的“升级版”，在不想牺牲可调试性、又希望减少 tracing 损耗时使用。

CUDA Graphs后端：一键播放的录像大师

有一类客人对出餐速度极其敏感——比如外卖超时就要退款。这时候小美会启用录像大师 cudagraphs。这位师傅的绝活是：第一次做这道菜时，把颠勺、开火、装盘的全过程录成高清视频；之后每次接单，直接按播放键，中间不需要再喊“开火！”“翻面！”——所有 GPU 内核的启动指令被预先录制在一张 CUDA 图里，执行时整张图一次性交给 GPU 回放。这就把 m 次内核启动的“口舌开销”压缩成了 1 次图播放。

$$\displaystyle T_{\text{graph}} = T_{\text{capture}} + T_{\text{launch\_graph}} + T_{\text{compute}}$$

1. 录制成本：$T_{\text{capture}}$ 是首次捕获 CUDA 图的一次性开销，录制完成后可反复播放，类似拍视频前期的布景。

$$T_{\text{capture}}$$

2. 单次播放启动：$T_{\text{launch\_graph}}$ 是播放整张图的启动开销，一次调用替代了原来 m 次内核启动，是延迟降低的核心。

$$T_{\text{launch\_graph}}$$

3. 计算主体：$T_{\text{compute}}$ 是所有内核实际执行的时间，图回放不会改变计算本身，只优化调度环节。

$$T_{\text{compute}}$$

不过录像大师有个死穴：他录的是“固定剧本”——输入尺寸一旦变化，录像就得重拍。所以 cudagraphs 最适合输入 shape 稳定、反复执行同一模型的在线推理场景，比如固定尺寸的 OCR 或图像分类服务。

后端选型策略：如何为你的菜品挑厨子

小美总结了一套“选厨心法”。第一步看模型结构：如果输入 shape 和网络分支完全固定，录像大师 cudagraphs 的杀手锏能发挥到极致；如果模型里充满动态控制流、shape 随样本变化，录像就容易报废，还是老实派更稳妥。第二步看硬件环境：inductor 的 Triton 内核在 NVIDIA GPU 上优化充分，而 CPU 或小众硬件上它的优势可能大打折扣。第三步算编译时间这笔账：开发调试期代码天天改，每次触发重编译都让人抓狂，eager 的白嫖优势就体现出来了。最后看迭代时长：单次训练上万步，务必选融合能力最强的 inductor，让编译成本被海量迭代摊薄。

$$\displaystyle R = \frac{\left( T_{\text{eager}} - T_{\text{backend}} \right) \times N}{T_{\text{compile}}}$$

1. 每步节省：$T_{\text{eager}} - T_{\text{backend}}$ 是相对 eager 基线，新后端每步省下的时间，是收益的来源。

$$T_{\text{eager}} - T_{\text{backend}}$$

2. 总收益：$\left( T_{\text{eager}} - T_{\text{backend}} \right) \times N$ 把单步收益乘上总迭代次数，得到整个任务的累计收益。

$$\left( T_{\text{eager}} - T_{\text{backend}} \right) \times N$$

3. 投资性价比：$T_{\text{compile}}$ 作分母，体现“研究菜谱”的前期投入。R>1 才值得换，R<1 说明大厨的备课时间比省下的还多。

$$\frac{\left( T_{\text{eager}} - T_{\text{backend}} \right) \times N}{T_{\text{compile}}}$$

小美把这条公式设为餐厅的经营守则：不是最贵的厨子最好，而是最匹配当前单子的厨子最好。先算账，再点将，这就是后端选型的第一性原理。

AI 应用场景

Python 演示

import torch  class SimpleMLP(torch.nn.Module):    def __init__(self):        super().__init__()        self.fc1 = torch.nn.Linear(512, 256)        self.fc2 = torch.nn.Linear(256, 128)        self.fc3 = torch.nn.Linear(128, 10)     def forward(self, x):        x = torch.relu(self.fc1(x))        x = torch.relu(self.fc2(x))        return self.fc3(x)  def run_with_backend(backend: str, iters: int = 30) -> float:    torch.manual_seed(0)    model = SimpleMLP().cuda().eval()    model = torch.compile(model, backend=backend)    x = torch.randn(64, 512, device="cuda")    with torch.no_grad():        for _ in range(5):            model(x)        torch.cuda.synchronize()        start_event = torch.cuda.Event(enable_timing=True)        end_event = torch.cuda.Event(enable_timing=True)        start_event.record()        for _ in range(iters):            model(x)        end_event.record()        torch.cuda.synchronize()        return start_event.elapsed_time(end_event) / iters  if __name__ == "__main__":    backends = ["inductor", "eager", "aot_eager", "cudagraphs"]    print("torch.compile 后端切换演示")    print("模型: SimpleMLP  GPU: " + torch.cuda.get_device_name(0))    for name in backends:        try:            avg_ms = run_with_backend(name)            print(f"{name:>12}  平均耗时 {avg_ms:.3f} ms/步")        except Exception as exc:            print(f"{name:>12}  运行失败: {exc}")

1. model = torch.compile(model, backend=backend)一行切换编译后端，模型代码完全不用改。
2. backends = ["inductor", "eager", "aot_eager", "cudagraphs"]枚举四种待测后端，循环逐一验证。
3. start_event.record()创建一个 CUDA 计时起点，记录 GPU 侧时间而非 Python 时间。
4. torch.cuda.synchronize()阻塞 CPU 直到 GPU 任务完成，确保计时覆盖全部内核执行。
5. start_event.elapsed_time(end_event) / iters用结束事件减去开始事件再除以迭代数，得到单步平均毫秒耗时。
6. try/except某些后端的 cudagraphs 模式在特定 GPU 或动态环境下可能不可用，容错保证循环继续。

常见误区

面试题