autograd（上） · 前向如何悄悄建图

autograd（上） · 前向如何悄悄建图

到了整台引擎最神奇的地方。你在 nanoGPT 里从没手写过一行求导，可 loss.backward() 一调，每个参数的 .grad 就有了。这不是魔法，是 autograd 引擎在你前向计算的同时，悄悄建了一张图。这一章讲「建图」，下一章讲「回放」。

核心思想：算的时候，顺手记账

autograd 的主意特别朴素：

你正常地做前向计算；每做一步运算，PyTorch 就顺手记一笔账——记下「这一步是什么运算、将来要对谁求导、需要留着哪些中间值」。前向走完，这些账就自动连成了一张「反向计算图」。

这叫 define-by-run（边跑边定义）/ 动态图：图不是你提前画好的，而是代码实际跑到哪、图就长到哪。有 if、有循环？那这次实际走了哪条分支，图就长成哪样。这也是 PyTorch 用起来像「普通 Python」的原因——你只管算，图在背后自己长。

三个关键词：requires_grad、grad_fn、saved tensors

拆开看这笔「账」记的是什么。跑一小段：

import torch
x = torch.tensor(3.0, requires_grad=True)   # 叶子,要求梯度
a = x * 2                                    # a.grad_fn = MulBackward
b = a + 1                                    # b.grad_fn = AddBackward
loss = b ** 2                                # loss.grad_fn = PowBackward

requires_grad：一个张量标了 requires_grad=True，就进入了「被追踪」名单。任何涉及它的运算，产出的结果也自动被追踪。 于是 x 一被追踪，a、b、loss 顺藤摸瓜全被追踪。

grad_fn：每个「算出来的」张量，都带一个 .grad_fn——指向一个反向节点(backward Node)，它知道这一步运算该怎么求导。a = x * 2 之后，a.grad_fn 是个 MulBackward；b = a + 1 之后，b.grad_fn 是 AddBackward。打印出来一看便知：

>>> a.grad_fn
<MulBackward0 object at ...>
>>> loss.grad_fn
<PowBackward0 object at ...>

saved tensors（留存的中间值）：反向求导常常需要前向的某些值。比如乘法 a = x * y，反向时 $\partial a/\partial x = y$、$\partial a/\partial y = x$——所以 MulBackward 必须把 x 和 y 留着，否则反向时算不出来。这些被留存的张量，就挂在对应的 grad_fn 上。

图是怎么连起来的

关键在于：每个 grad_fn 节点，除了「怎么求导」，还记着一组 next_functions——指向它每个输入的 grad_fn 的边。于是前向每算一步，就往图里加一个节点、连几条指回上游的边。前向走完，一张逆向的有向图就成型了：

注意箭头指向上游——这正是反向传播要走的方向。前向是 x → a → b → loss，而这张图是反着连的 loss → b → a → x。还记得上一章说 autograd 是 dispatcher 链上的一层拦截键吗？建图这件事，就发生在那一层：每个算子的 Autograd 键内核，一边 redispatch 去真正计算，一边顺手把这个反向节点建好、连好。

叶子、AccumulateGrad，与那些看不懂的报错

• 叶子张量(leaf)：你直接创建的、requires_grad=True 的张量（模型参数、输入）。它们是图的末梢，梯度最终要堆进它们的 .grad。图里那个特殊的 AccumulateGrad 节点，干的就是「把算到这儿的梯度累加进 x.grad」——nanoGPT 训练循环里 optimizer.zero_grad() 要清的，正是这个累加出来的 .grad。
• 非叶子（中间）张量：像 a、b，有 grad_fn 但默认没有 .grad（反向时梯度只是路过它们、传给上游，不停留）。

建图的代价：显存

这张图不是白建的，它有实打实的显存代价：那些 saved tensors（前向的中间激活），为了反向能用，必须一直留在显存里，直到反向用完才释放。所以：

• 训练为什么吃显存？ 很大一部分，就是这些为反向而留存的中间激活。模型越深、batch 越大，留存的越多。
• torch.no_grad() 为什么省显存？ 因为它关掉追踪——不建图、不留中间值。nanoGPT 的 generate 和 estimate_loss 都套了 @torch.no_grad()：推理和评估不需要反向，关掉建图，显存和速度都省一大截。
• 梯度检查点(gradient checkpointing) 这类技术，就是「反向时重新算」换「前向不留中间值」，拿计算换显存——本质也是在和这张图的存储代价做交易。

图建好了，该回放了

前向走完，一张连着 saved tensors、指向上游的反向图，已经静静躺在那儿。现在你调 loss.backward()——引擎会从 loss 出发，逆着这张图往回走，用链式法则一路把梯度算出来、堆进每个叶子的 .grad。这个「回放」是怎么跑的，下一章见分晓。👉 autograd（下）：backward 如何逆向回放

参考文献

• [1] Paszke et al. “Automatic differentiation in PyTorch.” NeurIPS Autodiff Workshop 2017. —— autograd 的设计。
• [2] PyTorch 文档：“Autograd mechanics.” https://pytorch.org/docs/stable/notes/autograd.html —— requires_grad、grad_fn、叶子、原地操作。
• [3] Edward Z. Yang. “PyTorch internals.” 2019. http://blog.ezyang.com/2019/05/pytorch-internals/ —— autograd 与 dispatch 的关系。
• 源码：torch/csrc/autograd/（尤其 function.h 的 Node）、tools/autograd/derivatives.yaml(pytorch/pytorch)。