autograd（下） · backward 如何逆向回放 · PyTorch 源码漫游

loss.backward() 之后，引擎从 loss 出发、按拓扑逆序回放计算图：每个节点拿到上游传来的梯度，用链式法则（向量-雅可比积）算出对输入的梯度，再往上游传；多条路径的梯度相加，最后堆进叶子的 .grad。看清 seed=1、梯度累加、图被释放这些细节。

autograd（下） · backward 如何逆向回放

上一章前向走完，一张反向计算图已经建好、静静躺着。这一章，我们调 loss.backward()，让 autograd 引擎逆着这张图跑一遍，看梯度怎么一步步算出来、堆进 .grad。

起点：给 loss 播一颗「种子梯度」

loss.backward() 一调，引擎就启动。第一件事，是给出发点一颗种子梯度：

loss 是标量时， $\partial \text{loss} / \partial \text{loss} = 1$ ，种子就是 1.0——所以标量 loss 直接 .backward() 就行。
loss 不是标量时，你得自己传 gradient=（那个向量），因为「对一个向量求导」得先说清「按什么方向」。这也点出了 autograd 的本质（下面讲）：它算的是向量-雅可比积，不是完整雅可比。

有了种子，引擎从 loss.grad_fn 出发，逆着图往回走。

每个节点做一件事：局部求导 × 上游梯度

引擎按拓扑逆序访问节点（保证轮到一个节点时，它下游传来的梯度都已备齐）。轮到某个节点（比如 MulBackward）时，它做三步：

收：拿到「输出的梯度」（从下游传上来的）；
算：用前向时 saved 的中间值，按这一步运算的局部导数，乘上收到的梯度——这就是链式法则；
传：把算出的「对各输入的梯度」，沿 next_functions 的边送给上游节点。

用乘法举例：a = x * y，前向留了 x、y。反向时它收到 ∂loss/∂a，于是：

\frac{\partial \text{loss}}{\partial x} = \frac{\partial \text{loss}}{\partial a}\cdot y, \qquad \frac{\partial \text{loss}}{\partial y} = \frac{\partial \text{loss}}{\partial a}\cdot x

把这两个梯度分别传给 x、y 的上游节点。整个 backward，就是这样一个节点接一个节点，把「局部导数 × 上游梯度」一路复合下去——链式法则在图上的忠实执行。

loss (种子=1)
  │  PowBackward: × 2b
  ▼
  b 的梯度 ── AddBackward: 原样传(加法导数=1)
  │
  ▼
  a 的梯度 ── MulBackward: × 2(因为 a = x*2)
  │
  ▼
  x 的梯度 ── AccumulateGrad: 堆进 x.grad

两条路汇一处：梯度要相加

一个张量若在前向里被用了多次（比如同一个 x 喂给了两个分支），反向时就会有多股梯度从不同路径涌回它。这时怎么办？——相加。这正是多元链式法则：一个变量通过多条路径影响最终结果，总梯度是各路径贡献之和。图里那个节点会等所有下游梯度到齐、加起来，再往上传。

叶子处的 AccumulateGrad 也是「加」：它把传到的梯度累加进 leaf.grad。

这就是为什么每步要 optimizer.zero_grad()。 .grad 是累加的，不是覆盖的。nanoGPT 训练循环里，一步更新完必须清零，否则这一步的梯度会叠在上一步没清的梯度上，越滚越错。而梯度累积那个技巧，反过来正是故意不清零、让多个 micro-batch 的梯度自然累加——用的就是 .grad 累加这个特性。原来两处都在利用同一个机制。

本质：反向模式自动微分，只算「向量-雅可比积」

这里点一个理论要害。一个函数从 $n$ 维输入到 $m$ 维输出，完整的雅可比矩阵是 $m\times n$ ——大得吓人。autograd 从不构造完整雅可比，它算的是向量-雅可比积(vector-Jacobian product， VJP)：给一个「输出方向向量」（就是那颗种子/上游梯度），直接算出「对输入的梯度」。

这就是反向模式自动微分的精髓：对于「多输入、单输出（标量 loss）」的场景——深度学习的典型情形——反向模式只需一遍反向就能拿到对所有参数的梯度，极其高效。每个节点的 backward 函数，本质就是「我这一步运算的 VJP」。autograd 引擎 = 把一堆局部 VJP，沿图复合起来。

收尾：图被释放，梯度就位

backward 跑完，默认(retain_graph=False)会释放这张图——那些为反向留存的 saved tensors 被放掉，显存回收。这就是「Trying to backward a second time」报错的来由：图已经拆了，你不能再回放一次（除非 retain_graph=True）。

此刻，每个参数的 .grad 都填好了。接下来 nanoGPT 里那句 optimizer.step() 就登场，拿着 .grad 去更新参数。你在 nanoGPT 训练循环读到的「前向 → 反向 → 更新」，底下正是：前向建图（上一章）→ 本章这台引擎回放 → 优化器用 .grad 走一步。 三章连起来，loss.backward() 的全部魔法就此揭穿。

想自己插一个「引擎不认识」的运算（比如调了个 C 库、或想手写更省的反向）？用 torch.autograd.Function：你实现 forward 和 backward 两个静态方法，PyTorch 就把它当成图里的一个节点——forward 里用 ctx.save_for_backward(...) 存中间值，backward 里读出来算 VJP。这是「往 autograd 图里手动加一个节点」的官方口子，也和 ATen 那本 derivatives.yaml 是一回事的两种形态（一个在 Python 手写，一个在 C++ 登记）。

autograd 讲完，回到上层和底层

引擎的心脏(dispatch + autograd)拆完了。接下来两章，一上一下：向上看你天天用的 nn.Module——它其实只是组织参数和 hook 的一层薄糖；向下看张量在 GPU 上怎么活、显存怎么被管理。先看上层。👉 nn.Module：参数、buffer 与 hook 的那层糖

参考文献

[1] Paszke et al. “Automatic differentiation in PyTorch.” NeurIPS Autodiff Workshop 2017. —— 反向引擎设计。
[2] PyTorch 文档：“Autograd mechanics” 与 “Extending PyTorch: torch.autograd.Function.” https://pytorch.org/docs/stable/notes/extending.html
[3] Baydin et al. “Automatic Differentiation in Machine Learning: a Survey.” JMLR 2018. arXiv:1502.05767. —— 反向模式 AD 与 VJP 的理论。
源码：torch/csrc/autograd/engine.cpp（反向引擎）、torch/autograd/function.py(pytorch/pytorch)。

autograd（下） · backward 如何逆向回放

本页目录