Lesson 0003：Backend 族谱 + 训练反传如何被编译

前两课你把 torch.compile 的流水线搞清楚了。这一课深入一个前两课刻意略过的基础—— 自动微分（autograd）到底怎么工作的， 然后再看 AOTAutograd 如何把它从"运行时查 tape"变成"编译时静态图"。 最后回答那个最关键的问题：TVM 到底能不能只复用 AOTAutograd 的联合图来做训练编译？

1. 先修：Autograd 基础 —— tape 是什么，梯度怎么算

PyTorch 的 autograd 是反向模式自动微分（reverse-mode automatic differentiation）的实现。 它是整个训练编译栈的基石——不理解它，就不理解为什么 AOTAutograd 是必要的。

1.1 从导数到梯度

我们先从最简单的概念开始，再逐步加入多变量和向量：

1.2 autograd 如何自动算梯度：计算图 + tape

PyTorch 的 autograd 分两步工作：

1.3 亲手看 tape：一个可跑的示例

import torch

x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)
b = torch.tensor([1.0], requires_grad=True)

# 前向 —— 边算边往 tape 上记录
y = x * w          # tape 记录 mul
z = y + b          # tape 记录 add
loss = z.sum()     # tape 记录 sum
print(f"loss = {loss.item()}")   # loss = 7.0  (x*w + b = 2*3 + 1)

# 反向 —— 逆序遍历 tape
loss.backward()

# 检查梯度
print(f"grad_x = {x.grad}")   # ∂loss/∂x = w = 3.0
print(f"grad_w = {w.grad}")   # ∂loss/∂w = x = 2.0
print(f"grad_b = {b.grad}")   # ∂loss/∂b = 1.0

你可以用 torchviz 或 torch.fx 把这张图可视化出来：

# 查看 autograd 图（需要 pip install torchviz）
from torchviz import make_dot
make_dot(loss, params={"x": x, "w": w, "b": b}).render("autograd_graph", format="png")

2. 编译栈的三层分工

在 autograd 的基础上，torch.compile 把编译栈分成三层。 不是每个 backend 都占三层——有些只占一层。

3. AOTAutograd：从动态 tape 到静态联合图

§1 讲的原生 autograd 有一个根本问题：tape 是运行时动态的。 每一次 forward 执行，tape 的结构可能不同（取决于数据流过的分支）。 编译器需要一张静态图才能做优化——所以在 Dynamo 抓到 forward FX Graph 之后， AOTAutograd 负责把 backward 也变成一张图。

3.1 AOTAutograd 的工作流

步骤②和③是 AOTAutograd 区别于 TVM 的核心能力： 在执行 forward 时，tape 被自动记录（这是 PyTorch C++ 核心做的事），然后 AOTAutograd 立刻读取 tape 上的信息，在 tape 被释放前把 backward 路径 trace 成 FX Graph。

3.2 联合图长什么样？一个可跑的具体例子

以下是一个最小但实际可跑的训练示例——注意 backward() 不在 torch.compile 包裹的函数里， Dynamo 不会去 trace 它。

import torch

# ── 定义纯前向模型（torch.compile 只包前向）──
def model(x, w, b):
    y = x @ w          # matmul
    y = y + b          # add
    y = y.relu()       # relu
    return y

# torch.compile 编译 model
compiled_model = torch.compile(model, fullgraph=True)

# ── 训练数据 ──
x = torch.randn(4, 4, device='cuda', requires_grad=True)
w = torch.randn(4, 4, device='cuda', requires_grad=True)
b = torch.randn(4, device='cuda', requires_grad=True)
target = torch.randn(4, 4, device='cuda')

# ── 训练步（backward 在 compile 之外）──
y = compiled_model(x, w, b)          # ← Dynamo 抓前向图
loss = (y - target).pow(2).sum()     # MSE loss（也在 compile 外）
print(f"loss: {loss.item():.4f}")

# backward 触发 AOTAutograd：
#   Dynamo 已经把 model 的 forward 抓成 FX Graph
#   → AOTAutograd 从 tape trace backward
#   → 拼接联合图 → Inductor 编译 → 执行反传
loss.backward()                       # ← 不经过 Dynamo，直接走 autograd 引擎
print(f"grad_w[0,0]: {w.grad[0,0]:.6f}")

4. Inductor 是谁开发的

Inductor 是 Meta（Facebook）PyTorch 团队开发的编译器后端， 作为 torch.compile 的默认 backend 随 PyTorch 2.0（2023 年 3 月）一起发布。

Inductor 的核心设计思路：不直接生成 CUDA 代码，而是生成 Triton 代码（GPU 路径） 或 C++/OpenMP 代码（CPU 路径）。之所以选 Triton 而不是 CUDA 作为代码生成目标， 是因为 Triton 的 block-level 抽象大幅降低了代码生成器的复杂度——这点在 §6 详讲。

5. 所有 Backend 的一张表

6. TVM 能不能只复用 AOTAutograd 做步骤⑤联合图编译？

这是整个课程最核心的技术判断。从 §3 的流程图看，AOTAutograd 的步骤⑤产出的是 FX Graph 格式的联合图——它是一张由 PyTorch ATen 算子组成的静态计算图。 TVM 的 Relay/Relax 前端确实可以接收 FX Graph 并转换为 Relay IR。

所以在纯理论上：AOTAutograd 产出联合图 → 转 Relay → TVM 编译 → 训练可行。

但工程实践上，这条路有五个死结：

死结 1：算子覆盖不完整

死结 2：动态 shape

TVM 的 auto-scheduling（AutoTVM / AutoScheduler）假设静态 shape—— 它针对固定的张量维度搜索最优 tiling 和 schedule 配置。但训练中 shape 经常是动态的：

死结 3：AOTAutograd 不是稳定公开 API

死结 4：In-place 操作和 optimizer 状态

训练不只是 forward + backward——还有 optimizer.step()：

死结 5：无人维护这条集成路径

这是最现实的死结。PyTorch 团队不维护、TVM 社区不维护、没有公司资助。 历史上存在过 torch_tvm 之类的实验项目（把 FX Graph 转 Relay）， 但它们都是推理 only，且大多已停止更新。

7. Triton 为什么能帮你写 kernel？

Inductor 选择 Triton 而非直接生成 CUDA 代码，是经过深思熟虑的架构决策。 要理解这个决策，先看两种编程模型的差异。

7.1 CUDA 的编程模型：thread-level

7.2 Triton 的编程模型：block-level

8. 反向传播是怎么做出来的？

从 torch.compile 的角度看，反向传播不是神秘的黑箱——它是一串确定的元素级操作， 可以被编译、可以被融合。

8.1 链式法则 = 本地乘 + 往前传

每个反传步骤只需要：前向时保存的中间值 + 上游传来的梯度 + 本层操作的局部导数公式。

8.2 常见算子的 backward

8.3 为什么编译能加速训练（不只是 forward）

同样的融合逻辑，forward 里有用，backward 里效果甚至更好——因为 backward 里充满了 element-wise 操作（梯度相乘、relu mask 应用、broadcast）， 这些是最容易融合的算子类型。

9. 动手验证：一条命令看到 forward+backward 的联合编译

在你 GPU 环境里跑这段——它能正常运行，不会报 Unsupported Tensor.backward()：

TORCH_LOGS="output_code" python -c "
import torch

# ═══ torch.compile 只包裹纯前向函数 ═══
def forward(x, w, b):
    y = x @ w          # matmul
    y = y + b          # add
    y = y.relu()       # relu
    return y

fn = torch.compile(forward, fullgraph=True)

# ═══ 数据 ═══
x = torch.randn(4, 4, device='cuda', requires_grad=True)
w = torch.randn(4, 4, device='cuda', requires_grad=True)
b = torch.randn(4, device='cuda', requires_grad=True)
target = torch.randn(4, 4, device='cuda')

# ═══ 训练步 ═══
y = fn(x, w, b)                        # ← Dynamo 抓前向图
loss = (y - target).pow(2).sum()       # ← MSE loss（不在 compile 里）
loss.backward()                        # ← AOTAutograd 首次 backward 时编译联合图

print(f'loss: {loss.item():.4f}, grad_w[0,0]: {w.grad[0,0]:.6f}')
" 2>&1 | grep -E "(triton_|def triton|backward|AOTAutograd)" | head -15

10. 小结

1. Autograd tape 是基础：前向记录操作序列，反向逆序回放。这是 PyTorch define-by-run 的核心——图不是在运行前声明的，而是在运行时动态构建的。
2. AOTAutograd 把动态 tape 变成静态联合图：前向跑一遍 → 读 tape → trace backward 图 → 拼接 forward+backward 联合图 → 丢给 Inductor 编译。它在第一次 backward 时触发，不在编译时。
3. Inductor 是 Meta 开发的默认 backend：随 PyTorch 2.0（2023.3）发布， 使用自研 Loop IR 做循环优化，生成 Triton（GPU）或 C++/OpenMP（CPU）代码。命名来自电感器（inductor）的工程隐喻。
4. TVM 理论上可复用 AOTAutograd，实践上不行：五个死结——算子覆盖不全、动态 shape 支持差、AOTAutograd 不是稳定 API、in-place optimizer 操作、无人维护集成路径。工作量相当于在 TVM 里重写半个 PyTorch。
5. Triton 抽象到 block 层：免去 thread 映射、shared memory 管理、同步这些 CUDA 的手工活——Inductor 不必操心这些就能生成高效 kernel。
6. Backward 也是 matrix ops + element-wise：编译一样有效，甚至更有效（backward 中 element-wise 占比更高）。
7. torch.compile 不包 backward()：只包裹前向函数，backward 在外层调用——AOTAutograd 在首次 backward 时自动介入。

11. Quiz

@torch.compile
def train_step(model, x, target):
    loss = loss_fn(model(x), target)
    loss.backward()  # ← 这行
    return loss