跳到主要内容

数据截至 (上游 commit c187ef3271d5)

03 · Autograd 引擎：反向图怎么挂、怎么跑

这一章讲什么： loss.backward() 背后那台机器。读完你会知道：反向图是什么时候、由谁挂上的；引擎为什么不显式做拓扑排序也能按正确顺序执行；梯度最后是怎么落到 param.grad 的。

先给一句定位：如果你只想理解「反向传播这个数学概念」，去读教学实现 micrograd（两百行 Python，递归 + 显式拓扑排序）。本章回答的是另一个问题：同一件事在生产规模下怎么做——多线程、多设备、每个算子的 backward 都是真 kernel。

1. 它要解决的小问题​

反向传播的数学很简单（链式法则），工程上难的是：

  • 图是运行时动态长出来的——PyTorch 是 eager 框架，没有静态图可分析；
  • 算子有上千个，每个的 backward 公式都要与前向 kernel 严格对应；
  • 前向可能跨设备、跨线程，反向执行要按依赖顺序调度，还要处理「一个节点有多个下游」「某个叶子不需要梯度」等现实情况。

PyTorch 的答案分两半：挂图（前向顺手做）和执行（一个专门的多线程引擎）。

2. 挂图：每个输出记住「我从哪来」​

2.1 图的两个基本件​

概念类型干什么位置
节点struct Node一个算子的 backward：operator() 吃进上游梯度、吐出对输入的梯度；next_edges_ 指回它前向时的各个输入torch/csrc/autograd/node.h:124、:166
边Edge = (Node*, input_nr)；edge_list「我的第 k 个输入来自哪个节点的哪个输出」torch/csrc/autograd/node.h:37
元数据struct AutogradMeta挂在 TensorImpl 上：grad_fn_（非叶子）、grad_accumulator_（叶子）、grad_、requires_grad_、is_view_torch/csrc/autograd/variable.h:231

判断一个 Tensor 在图里的身份只看一个字段：

  • grad_fn_ 非空 → 中间节点（非叶子），反向时执行它；
  • is_leaf 且 requires_grad_ → 叶子，反向时把梯度累加进 .grad。

2.2 谁挂的图：生成代码 VariableType​

第 2 章说过 dispatch 会经过 autograd 层。这一层的代码不是手写的，是构建期生成的：tools/autograd/gen_variable_type.py 为每个算子生成 VariableType 包装，核心动作就是给输出调用 rebase_history（代码模板见 tools/autograd/gen_variable_type.py:716，底层符号在 torch/csrc/autograd/variable.h:170）——把「本算子的 Node + 指向各输入的边」记到输出的 grad_fn_ 上。

每个算子的 backward 公式则写在一张表里：tools/autograd/derivatives.yaml。以 add.Tensor 为例（tools/autograd/derivatives.yaml:229-232）：

- name: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  self: handle_r_to_c(self.scalar_type(), grad)
  other: handle_r_to_c(other.scalar_type(), maybe_multiply(grad, alpha.conj()))
  result: self_t + maybe_multiply(other_t, alpha)

即：加法对两个输入的梯度都是「原样传回上游梯度」（复数情形共轭）。这张 yaml 是数学和工程的接缝——改错一行，反向就悄悄算错。

2.3 直觉演示​

# 示意，非源码
x = torch.tensor(2.0, requires_grad=True)   # 叶子：grad_fn=None
y = x * x       # y.grad_fn = MulBackward，next_edges → [(AccumulateGrad(x), 0)]
z = y + y       # z.grad_fn = AddBackward，两条边都指回 y 的 grad_fn
z.backward()    # 从 z 的 grad_fn 出发，倒着走到 AccumulateGrad(x)

这张图就是一个 DAG：节点是 backward 函数，边是「梯度流向」。注意它不存前向数值本身——节点只按需要存了 backward 要用的东西（比如 MulBackward 会存下输入 x）。

3. 执行：Engine 的拓扑调度​

3.1 入口​

tensor.backward()（torch/_tensor.py:566）→ torch.autograd.backward（torch/autograd/__init__.py:255）→ _engine_run_backward（torch/autograd/graph.py:1083）→ 最终调进 C++：Variable._execution_engine.run_backward（torch/autograd/graph.py:1096）→ Engine::execute（torch/csrc/autograd/engine.cpp:1363）。

3.2 核心机制：依赖计数，而不是拓扑排序​

micrograd 的做法是先 DFS 出拓扑序再倒序执行。PyTorch 的做法更并发友好：

  1. 建 GraphTask 时给每个节点数入度——它有几个「下游」要走完才能执行它（dependencies_，在 evaluate_function 里查，torch/csrc/autograd/engine.cpp:1224）。
  2. 每执行完一个节点，把它每个 next_edge 指向的节点依赖数减一（--it->second == 0 即就绪，torch/csrc/autograd/engine.cpp:1232-1234）。
  3. 就绪就推进 ReadyQueue，由 worker 线程取走执行；未就绪的暂存在 not_ready_ 缓冲里等剩下的输入（torch/csrc/autograd/engine.cpp:1235-1280）。

调度循环的骨架（示意，非源码）：

# 示意，非源码
while queue.not_empty():
    node, input_buffer = queue.pop()
    grads = node(input_buffer)            # 执行该算子的 backward
    for edge in node.next_edges:          # 把梯度发给各输入节点
        next_node = edge.function
        accumulate_into(input_buffers[next_node], edge.input_nr, grad)
        if --dependencies[next_node] == 0:
            queue.push(next_node)         # 依赖清零，可以执行了

真实代码里这一步在 Engine::evaluate_function 的后半段（torch/csrc/autograd/engine.cpp:1103 起，梯度分发在 :1218-1286）。线程模型是：CPU 一个 ready queue，每类设备有自己的 worker 线程（Engine::thread_init，torch/csrc/autograd/engine.cpp:340），NodeTask 按设备路由到对应队列（ready_queue(cpu_ready_queue, next.function->device())，torch/csrc/autograd/engine.cpp:1254）。

3.3 终点：AccumulateGrad​

叶子参数也有一个 Node，叫 AccumulateGrad（torch/csrc/autograd/functions/accumulate_grad.h:43）。它的 apply（torch/csrc/autograd/functions/accumulate_grad.cpp:92）做的事就一件：把传来的梯度加进 param.grad（不是赋值——所以多个分支的梯度会在叶子处汇合，所以要手动 zero_grad）。

3.4 一个值得记住的警告​

backward(create_graph=True) 会在参数和它的梯度之间造成引用环导致泄漏，引擎在 Engine::execute 里直接打了 TORCH_WARN_ONCE（torch/csrc/autograd/engine.cpp:1375-1382），并建议改用 autograd.grad。这类「官方在代码里劝你」的注释往往比文档更诚实。

4. 自定义算子怎么进图：torch.autograd.Function​

Python 侧给用户的扩展点是 Function（torch/autograd/function.py:560）：你写 forward/backward 两个静态方法，用 ctx.save_for_backward（torch/autograd/function.py:41）存前向张量，调用 MyFunc.apply(x) 时框架会包出一个 Node 挂进同一张图。

# 示意，非源码
class Square(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)      # 存 backward 要用的东西
        return x * x
    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return 2 * x * grad_out       # 对输入的梯度

重点看：它和 ATen 内建算子在图里地位相同——引擎根本不知道一个 Node 背后是生成代码还是你的 Python 函数。

5. 坑与边界​

  • 图默认用一次就释放。 二次 backward() 报 "Trying to backward through the graph a second time"；要保留得 retain_graph=True。
  • in-place 与 view 的冲突检测靠 TensorImpl 上的 version counter；报错文案（"a view of a leaf ..." 之类）看着玄，本质就是版本号对不上。
  • .grad 是累加的。 忘了 zero_grad 梯度就越滚越大；Optimizer.zero_grad 现在默认 set_to_none=True（torch/optim/optimizer.py:1058）。
  • 多线程反向有讲究：reentrant backward（backward 里再 backward）有专门的工作窃取机制，见 torch/csrc/autograd/engine.cpp:144-167 的大段注释。
  • 看不出来：evaluate_function 里 CUDA stream/event 同步的全部边角（torch/csrc/autograd/engine.cpp:1108-1138 有 stream guard 与 event wait 逻辑），跨流反向的正确性细节建议直接读该函数。

6. 本章代码地图​

主题文件路径符号名
张量侧的梯度元数据torch/csrc/autograd/variable.hAutogradMeta、set_gradient_edge、rebase_history
反向图节点torch/csrc/autograd/node.hNode、Node::operator()、Edge
引擎主入口torch/csrc/autograd/engine.cppEngine::execute
节点求值与分发torch/csrc/autograd/engine.cppEngine::evaluate_function
就绪队列torch/csrc/autograd/engine.cppReadyQueue::push/pop
叶子累加torch/csrc/autograd/functions/accumulate_grad.hAccumulateGrad
求导公式表tools/autograd/derivatives.yamladd.Tensor 条目
挂图层代码生成tools/autograd/gen_variable_type.pyrebase_history 模板
Python 入口torch/autograd/__init__.py、torch/autograd/graph.pybackward、_engine_run_backward
自定义 Functiontorch/autograd/function.pyFunction、FunctionCtx.save_for_backward