软件基座：框架与自动微分

在模型的数学与 第 62 章 的加速器之间，还有一层本书到目前为止一直默认存在的东西。第一部分的每一份训练配方都假定：一行代码 loss.backward()，就能算出标量损失对几十亿个参数的梯度，而且算得很便宜。这个假定背后压着一条定理、一种实现技术，以及框架设计围绕如何暴露它们而争了二十年的历史。本章把三者依次展开：为什么一个程序的梯度只花费程序本身的一个小常数倍，框架如何记录并回放计算来收集梯度，计算图编译器与即时执行之间长达十年的争论如何收场，以及当把模型切分到集群成为编程接口的一部分之后，框架变成了什么。读完本章，读者能说清从写下损失函数到硬件底座开始做算术之间到底发生了什么，以及这条路径为何长成这样。

支撑整个技术栈的一行代码

第一部分讲过，训练就是梯度下降：算损失，算损失对每个参数的梯度，往下坡走一步，重复。损失容易，它只是模型的一次前向运行。需要想法的是梯度，因为对「给模型求导」的朴素理解会指向两条路，而两条在这个规模上都走不通。

第一条是数值差分：把某个参数拨动一个小量 $h$，重新跑一遍模型，用损失的变化除以 $h$。这样每个参数要花一次前向运行，十亿参数就是每步十亿次前向；而且答案两头受损，$h$ 太大有截断误差，$h$ 太小有浮点舍入误差 (Baydin et al. 2018)。第二条是符号微分，也就是计算机代数的做法：把损失当作数学表达式来变形，导出一条导数公式。可是复合表达式的导数公式可能比表达式本身大出指数倍，这种失败模式叫表达式膨胀（expression swell）；何况公式跟不上真实程序里的分支与循环 (Baydin et al. 2018)。

自动微分（automatic differentiation）走的是第三条路。它不对公式求导，而对执行过程求导：让程序照常运行，在每个基本运算处，一次乘法、一次加法、一次 sin，就对流过的数值应用这一个运算的求导法则。链式法则随后把这些局部导数组合成程序实际所做之事的精确梯度，分支与循环全都包含在内，精确到机器精度。没有步长要调，也没有表达式会膨胀。剩下的问题只有成本，而成本的答案，正是这一层的奠基定理。

反向模式：一次回扫换来全部梯度

把程序的执行写成一列基本步骤，每一步从更早的值算出一个中间值 $v_i$。这份线性化的轨迹有个比这个领域还老的名字：Wengert 列表，得名于 1964 年那篇两页纸的论文 (Wengert 1964)。于是求导有了方向。前向模式让导数伴着数值一起往前推，回答的是「这个输入动一动，下游的一切怎么动」；一趟只能得到对一个输入的敏感度，所以标量损失对 $n$ 个参数的梯度要跑 $n$ 趟。反向模式则先把轨迹正向跑一遍、存下中间值，再从输出往回扫，对每个中间值问「最终的损失有多在乎你」，并把答案累积起来，习惯上写作 $\bar v_i = \partial L / \partial v_i$，称作 $v_i$ 的伴随量。一次前向加一次回扫，就同时得到单个输出对所有输入的敏感度 (Baydin et al. 2018)。

对训练来说，这不是一道需要犹豫的选择题。损失只有一个标量，参数却数以十亿计。反向模式用一次前向加一次回扫拿到整条梯度，而廉价梯度原理把上界钉得很死：只数乘法，梯度的成本至多是函数求值的 3 倍，这是 Baur 与 Strassen 的结果 (Baur and Strassen 1983)；数上全部运算，保守上界是 5 倍，与函数有多少个输入无关 (Griewank and Walther 2008; Griewank 2012)。反向传播，那个为分层网络重新发现、在 1986 年声名大噪的算法 (Rumelhart et al. 1986)，正是反向模式在神经网络上的特化；一般形式早在十年前就已为数值程序发表 (Linnainmaa 1976; Speelpenning 1980)。值得停下来体会这个结果有多奇怪：对一个输入的敏感度，与对十亿个输入的敏感度，花的都是前向运行的同一个小倍数。完整的雅可比矩阵，即每个输出对每个输入的敏感度矩阵，没有任何类似的结果；这条定理是损失函数「多输入、单输出」形状的特有性质，而深度学习的整套经济学就压在它上面。

下面的可运行单元是一台完整的反向模式引擎，只有十五行上下，而且原理上并不是玩具：PyTorch 的 autograd 就是这一设计配上工业级的算子集合。每个算术运算记下它的父节点，以及输出对每个父节点的局部导数；这份记录就是 Wengert 列表，由普通的运算符重载在前向运行时顺手建成。对输出调用 backward()，就按逆序走一遍记录，把伴随量相乘再累加。注意 x 喂给了两个运算，它的梯度分两份到达、必须相加，这正是各家的「磁带」（tape，按执行顺序记下的运算列表）都往 .grad 里累加而不是赋值的原因，也是 PyTorch 程序每步都要调 zero_grad() 的原因。

import math

class Var:
    def __init__(self, val, parents=()):
        self.val, self.parents, self.grad = val, parents, 0.0
    def __add__(self, o): return Var(self.val + o.val, [(self, 1.0), (o, 1.0)])
    def __mul__(self, o): return Var(self.val * o.val, [(self, o.val), (o, self.val)])
    def backward(self, seed=1.0):           # 逆序走磁带
        self.grad += seed                   # += ：被用两次的值收到两份贡献
        for parent, local in self.parents:
            parent.backward(seed * local)

def sin(v): return Var(math.sin(v.val), [(v, math.cos(v.val))])

x, y = Var(2.0), Var(3.0)
z = x * y + sin(x)          # 前向传播顺手把磁带记了下来
z.backward()                # 一次回扫，拿到每个输入的梯度
print(x.grad, "== y + cos(x) =", 3.0 + math.cos(2.0))
print(y.grad, "== x")

这份草图与真实引擎之间隔着一个必须交代的差别：这个版本沿每条路径递归，一个可由多条路线到达的值会被访问多次。生产级的磁带会先做拓扑排序、每个节点只访问一次，但执行的算术是一样的。

代价是一条驻留内存的磁带

定理约束的是时间，账单却悄悄寄给了内存。回扫需要前向传播的中间值，也就是激活值，所以磁带的内存随计算长度增长。对一个深层网络，这意味着整个批次在每一层的激活值都要留到反向传播消费为止，训练一个模型比运行它多花好几倍内存，原因就在这里。

解药和这个领域一样老，而且仍在生产环境服役。Griewank 在 1992 年证明：沿轨迹布置检查点，回扫时重算检查点之间的片段，就能用对数倍的额外计算换来对数级的内存 (Griewank 1992)。深度学习对它的重新发现，亚线性内存训练，把这笔交换在层状网络上算得很具体：只在 $O(\sqrt{n})$ 的层上保留激活值，其余按需重算，一个千层网络的训练内存从 48 GB 降到 7 GB，时间多花约 30% (Chen et al. 2016)。框架把它作为激活重计算发布，而读者其实已经从硬件那一侧见过它的指纹了。

从计算图到磁带：框架之战

框架是上述机器的工业化成品，而这个领域花了十年才对它的形状达成一致。争的其实是一个设计决定：框架应该在什么时候看到程序？

第一个阵营的回答是「运行之前」。2010 年出自 Bengio 实验室的 Theano 让用户先搭一张符号表达式图，框架再对图做优化、把求导当作图变换完成、编译成 C++ 或 CUDA (Bergstra et al. 2010)。TensorFlow 在 2015 年把同一套「先定义后运行」（define-and-run）的设计推广到 Google 的机群规模：由会话执行一张静态数据流图，图可以切分到数百台机器上 (Abadi et al. 2016)；它的前身 DistBelief 早已在数万个 CPU 核上训练过十亿参数的网络，只是那更像一个被配置的系统，而不是被编程的系统 (Dean et al. 2012)。整图视角换来的东西是实打实的：框架可以融合算子、规划内存、序列化模型、把子图安放到设备上。付出的代价则是程序员本身。图程序是「构建程序的程序」，形状错误会在远离肇事代码行的地方爆出来，调试器无法单步进入尚未运行的东西，循环也必须写成图算子而不是一句 for。

第二个阵营的回答是「运行之时」。2015 年，Chainer 给另一条路起了名字：「运行即定义」（define-by-run），即时执行普通的宿主语言代码，边发生边记录运算，再对记录下来的磁带求导，于是网络结构就是 Python 实际做过之事的轨迹，分支、递归，全都算数 (Tokui et al. 2015)。Torch7 在 Lua 里培育同一种命令式文化已有多年 (Collobert et al. 2011)，而 HIPS 的 autograd 项目证明磁带可以贴着朴素的 NumPy 运行 (Maclaurin et al. 2015)。2017 年 1 月，PyTorch 把这几条血脉熔在一起：即时执行的磁带，底下是 Torch 的 C 库，工程上做到让易用性留在 Python 里，热路径却绕开 Python (Paszke et al. 2019)。

研究界用脚投了票，而且票数可以数出来：到 2019 年，在可识别框架的 CVPR 论文里 PyTorch 占 69%，在主要 NLP 会议上占绝对多数，TensorFlow 的研究份额则停滞或收缩 (He 2019)；2020 年 1 月，OpenAI 宣布全面标准化到 PyTorch (OpenAI 2020)。同一时期，图阵营在接口上认了输：TensorFlow 2.0 于 2019 年把即时执行设为默认。磁带赢下了「模型怎么写」的争论。至于它是否也赢下了「模型怎么跑」，那是下一章的问题，因为计算图的优化能力从来不是错的那一半。

2018 年 12 月由 Google 开源的 JAX，与其说是第三个阵营，不如说是两边的合题。它的程序看上去是即时的，就是数组上的朴素 NumPy；但它的变换，grad、jit、vmap，靠的是把一个纯函数追踪成图、交给 XLA 编译器 (Frostig et al. 2018)。合题的代价是纯函数性：副作用在追踪之下消失，依赖数据的 Python 控制流必须改写成结构化的组合子，这正是 TF1 当年要求过的让步，只是如今被限定在编译区域之内。HIPS autograd 的两支后裔就这样把设计空间一分为二：PyTorch 留住宿主语言、再补装编译器，JAX 留住编译器、去约束宿主语言。

合流，以及悬而未决的部分

把这段历史当成一个整体来读，这场战争以一笔交换收场。即时语义赢下编程模型，因为研究者调试的是程序，不是图；图执行赢下运行时，因为 第 62 章 的硬件底座奖励融合过、规划过、提前调度好的工作。每一个活下来的技术栈，如今都是两者之间的某种谈判。TensorFlow 2 即时运行，通过 tf.function 分阶段成图。JAX 在 jit 之下运行被追踪的纯函数。PyTorch 先试过静态语言子集这条路，即 TorchScript，官方复盘承认它只在大约一半的真实模型上可用，因为它要求整个程序都服从它；2023 年 3 月起稳定的 torch.compile 把赌注反了过来：能捕获的捕获，捕获不了的断图退回 Python，从构造上保住即时语义 (Ansel et al. 2024)。捕获如何进行、捕获到的图又被下面的编译器怎样处理，是 第 64 章 的主题。

合流没有解决的，是这笔交换买来了什么、又花掉了什么。即时执行的一步要为每个算子付一次 Python 分发的开销：当每个算子都是一次大矩阵乘时可以忍受，当内核只有微秒长时就是灾难。被捕获的图摊薄了分发开销，却必须察觉自身假设何时失效并重新编译，于是这笔交换又以编译延迟与缓存未命中的形式回来。本章开头的定理对这个选择无动于衷，常数因子却不然，而生产系统就活在常数因子里。

今天的框架是什么

剥掉品牌，2026 年的框架是叠在自动微分核心上的五项承诺，其中最后一项是新的。

第一项是张量库，带一套稳定的算子接口：几百个数学运算，形状与数据类型都有约定。第二项是自动微分引擎，工业级的 Wengert 列表。第三项是分发器，按设备、精度与模式把每个运算路由到具体内核，同一个程序能跑在 NVIDIA GPU、TPU 或一台笔记本上，靠的就是换掉栈底的这套机制。第四项是设备运行时，而它最容易被低估：用流让加速器与 Python 保持异步，用一个缓存式内存分配器回收显存，之所以要自己回收，是因为厂商自带的释放例程可能让整块设备同步 (Paszke et al. 2019)。每个走到生产环境的框架都长出过一个分配器；它在碎片化之下的行为，是 第 31 章 那个世界里反复出现的运维事故。

第五项承诺是集群。分布式过去不在框架里，而在包装库和启动脚本里；如今它以分片（sharding）的形式搬进了编程模型：用标注声明一个张量如何切分到设备网格上，由系统插入通信。GSPMD 在 XLA 内部演示了这个设计，寥寥几个标注被编译器传播到整张图 (Xu et al. 2021)；它如今的化身是 JAX 的命名分片与 shard_map，以及 PyTorch 在 DeviceMesh 之上的 DTensor，再往上是打包好的分片训练配方 FSDP (Zhao et al. 2023)。并行的数学属于 第 10 章；属于本章的是一个接口事实：一次训练运行的布局如今写在框架里，而不是写在包着框架的脚本里。正因如此，第 62 章 的带宽层级才会以类型签名而非部署细节的形式出现在程序员面前。

从制度上看，这一层已经按基础设施的方式稳定下来：PyTorch 在 2022 年从 Meta 移入 Linux 基金会旗下的实体，到 2025 年那里还接纳了 vLLM 与 DeepSpeed；JAX 仍是 Google 主导的开源项目，其编译器基座则由 OpenXLA 单独治理。到 2026 年年中，稳定版本线是 PyTorch 2.12 与 JAX 0.10，PyTorch 的分片张量机制已有文档但仍正式标注为 alpha，这恰好说明第五项承诺还有多年轻。

延伸阅读