第 2 讲 · ML 系统 101：基础 + TensorFlow / PyTorch · ML Systems Basics

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 2 · Wed 9/3 · Introduction to ML Systems。 阅读材料：（必读）Abadi et al., TensorFlow: A System for Large-Scale Machine Learning（OSDI’16）；（必读）Paszke et al., PyTorch: An Imperative Style, High-Performance Deep Learning Library（NeurIPS’19）。 依据说明：本讲以原课第 2 讲公开的 slides（02-MLSys-basics.pdf，47 页，“ML Systems 101 (Basics + TensorFlow/PyTorch)“）为主线。第一部分：深度神经网络 = 张量 + 算子；计算图级优化（Conv+BN 融合、两个 Conv 融合、自动发现图优化）；深度学习系统的组成；SGD 回顾；并行化（数据并行、模型并行、流水线并行、张量并行、组合与 3D 并行）；ML 编译（张量、张量函数、目标、厂商算子库与自动代码生成）；内存优化（重计算、零冗余、计算—内存—通信的平衡）。第二部分：TensorFlow 的延迟执行模型、PyTorch 的命令式模型、两者对比、PyTorch 2.0。§9 的论文精读依据两篇论文原文。minigraph.py 迷你系统、全部实验和练习是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture02/（4 个脚本 + 共享模块 minigraph.py；numpy + PyTorch CPU 版；每个脚本 ≤ 2 秒） 预计阅读时间：4–5 小时（含读两篇论文）

0. 本讲导航 · How to read this lecture

0.1 学习目标

学完这一讲，你应该能够：

1. 说出一个深度学习系统的”标准部件”：张量与算子、计算图、自动微分、计算图优化、并行化、ML 编译、内存优化，以及每个部件解决什么问题。
2. 解释反向模式自动微分为什么只需要每个算子提供一个”向量-雅可比积”（VJP），并区分两种实现：图上的符号微分（TensorFlow）和运行时的磁带（PyTorch）。
3. 推导 Conv + BatchNorm 折叠和两个卷积合并两条图改写规则，知道它们成立的前提条件，理解为什么”自动发现图优化”需要搜索和验证。
4. 写出数据并行、张量并行、流水线并行各自依赖的数学等式，以及流水线空泡比例 $(P-1)/(M+P-1)$。
5. 解释重计算为什么能把激活显存从 $O(N)$ 降到 $O(\sqrt N)$，代价是多少计算。
6. 比较 TensorFlow 的延迟执行和 PyTorch 的即时执行：各自让哪些事情变得容易、哪些变得困难；PyTorch 2 的 torch.compile 怎样兼得两者，又在哪里会”断图”。
7. 按第 1 讲的方法精读两篇必读论文，写出三段式评审。

0.2 先修与本讲范围

0.3 本讲地图

§1  从第 1 讲出发：一个深度学习系统由哪些部件组成
§2  张量、算子、计算图；SGD 的三步
§3  自动微分：VJP、反向模式、符号微分 vs 磁带                ← 代码 01（演示 1）
§4  计算图级优化：Conv+BN 折叠、卷积合并、"先变差再变好"      ← 代码 02
§5  并行化概览：数据、张量、流水线                          ← 代码 04
§6  ML 编译概览：张量函数、厂商算子库的局限、自动代码生成
§7  内存优化：重计算与零冗余                                ← 代码 03
§8  TensorFlow vs PyTorch：延迟执行 vs 即时执行；PyTorch 2    ← 代码 01（演示 2–4）
§9  论文精读：TensorFlow（OSDI'16）、PyTorch（NeurIPS'19）、两篇的联系
§10 小结   §11 常见误区   §12 练习   §13 延伸阅读   §14 下一讲预告

建议：先读 §1–§3，运行 01 的演示 1；读 §4 运行 02；§5–§7 运行 04、03；§8 回到 01 的演示 2–4；最后读两篇论文，再读 §9 对照。

1. 从第 1 讲出发：一个深度学习系统由哪些部件组成 · Anatomy of a DL system

第 1 讲把 ML 系统画成一个六层的栈。第 2 讲的 slides 开头重画了这张图，并在最上面加了一层：自动微分（Automatic Differentiation）。完整的部件清单（slides 的”深度学习系统概览”）是：

为什么要先把这张”零件清单”讲一遍？因为这门课后面的每一讲，几乎都是深挖其中一个零件。而且你会发现，一个系统怎样实现这些零件，很大程度上取决于一个更基本的设计选择：是先看到整个计算再执行，还是写一行执行一行？——这就是 TensorFlow 和 PyTorch 的根本分歧（§8）。

2. 张量、算子、计算图 · Tensors, operators, computation graphs

2.1 深度神经网络 = 张量 + 算子

slides 对深度神经网络的定义很朴素：一组简单的、可训练的数学单元，协作完成复杂任务。在系统的视角下，它只由两种东西组成：

• 张量（tensor）：$n$ 维数组。输入、输出、中间结果、参数都是张量。
• 张量代数算子（operator）：对张量做的运算，例如卷积、矩阵乘法、加法、ReLU、softmax。

把算子画成节点、张量画成边，就得到一张计算图（computation graph，也叫数据流图 dataflow graph）：边表示”这个算子的输出是那个算子的输入”。一个两层 MLP 的计算图：

x ──► matmul ──► add ──► relu ──► matmul ──► add ──► softmax_xent ──► loss
        ▲         ▲                 ▲         ▲            ▲
        W1        b1                W2        b2           y

为什么计算图这么重要？ 因为它是一个和执行无关的、可分析的程序表示。拿到它，系统就可以在执行之前回答很多问题：哪些算子可以并行？哪些中间结果什么时候可以释放？哪几个算子可以合成一个内核？这张图能不能切开放到多块 GPU 上？这些问题对应 §4–§7 的每一种优化。

2.2 训练：SGD 的三步

slides 回顾了随机梯度下降（SGD）的三个阶段，每一次迭代都要做一遍：

1. 前向传播：把一批输入样本送进模型，沿计算图逐个算子算出预测和损失；
2. 反向传播：沿计算图”倒着走”，算出损失对每个可训练参数的梯度；
3. 参数更新：$w \leftarrow w - \eta \nabla_w L(w) = w - \frac{\eta}{n}\sum_{i=1}^{n} \nabla_w \ell_i(w)$。

注意最后这个式子：一个批次的梯度是每个样本梯度的平均。这个看似平凡的事实，是数据并行（§5）能够成立的全部原因。

2.3 一个能亲手拆开的迷你系统：minigraph.py

为了能亲手比较两种执行方式，本讲的代码从零写了一个迷你深度学习系统 minigraph.py（约 900 行，大半是注释；NumPy 实现）。它的结构和真实框架一致：

• 算子表 OPS：每个算子登记四样东西——前向计算、形状推断、反向规则（VJP）、FLOPs。
• 一个入口 apply(op, *inputs)：输入是符号节点 Node 时，只往图里加一个节点、不计算（TensorFlow 1.x 风格）；输入是即时张量 Tensor 时，立刻算出结果，并在”磁带”上记一笔（PyTorch 风格）。
• 执行器 run(fetches, feed)：只执行 fetches 需要的子图，用引用计数释放中间结果，统计内核数和内存峰值。
• 图优化 pass：常量折叠、CSE、Conv+BN 折叠、卷积合并等。

第 10、11、15 讲还会继续用它。

3. 自动微分 · Automatic differentiation

3.1 直觉：链式法则 + 倒着走一遍

训练需要损失 $L$ 对每个参数的梯度。手推梯度对几十层的网络不现实，所以所有框架都提供自动微分（AD）：用户只写前向计算，系统自动算出梯度。

AD 的核心是链式法则。设计算图里某个算子 $y = f(x_1, \ldots, x_k)$，并且我们已经知道了损失对它输出的梯度 $\bar y = \partial L / \partial y$（这个记号读作”y bar”，叫做 $y$ 的伴随，adjoint）。那么损失对每个输入的梯度是：

当 $x_i$、$y$ 是张量时，$\partial f / \partial x_i$ 是一个巨大的雅可比矩阵，但我们从来不需要把它显式构造出来，只需要计算”伴随向量乘以雅可比矩阵”这个乘积——它叫做向量-雅可比积（vector-Jacobian product，VJP）。每个算子的 VJP 都有简单的闭式：

有了每个算子的 VJP，整个网络的梯度就是：从损失出发（$\bar L = 1$），按拓扑序的逆序遍历计算图，对每个节点用它的 VJP 把伴随传给输入，多条路径汇合处把伴随相加。 这就是反向模式自动微分。

3.2 为什么是”反向”模式？

AD 还有一种”前向模式”：从某个输入出发，顺着计算图把”对这个输入的导数”往前推。两者的代价：

• 前向模式：一次遍历得到所有输出对一个输入的导数；
• 反向模式：一次遍历得到一个输出对所有输入的导数。

训练的情形是”一个标量损失，几十亿个参数”，所以反向模式一次遍历就够了，代价只是前向的常数倍（一般 2 倍左右）。PyTorch 论文也提到了这一点：反向模式适合”标量输出、多元输入”，前向模式适合输出比输入多的情形，后者在机器学习里不常见。

反向模式的代价是内存：VJP 规则里要用到前向的中间结果（矩阵乘法的 VJP 要用 $A$ 和 $B$，tanh 的 VJP 要用输出 $y$），所以前向的激活要一直保留到反向用完为止。§7 的重计算就是为了缓解这个问题。

3.3 两种实现：符号微分（图）vs 磁带（即时）

同样的 VJP 规则，有两种组织方式：

• 符号微分（TensorFlow）：在计算图上，从损失节点开始按逆序遍历，把每个 VJP 也变成图里的节点。结果是一张更大的图：前向子图 + 反向子图。反向和前向一样，都是”可以被分析和优化的图”。
• 磁带（tape，PyTorch）：即时执行时，每算一个算子，就在”磁带”上记一笔”这个结果是哪个算子、用哪些输入算出来的”。调用 loss.backward() 时，沿磁带倒着回放，直接计算梯度。图是”边执行边记录”的，每次前向都重新记一遍——所以模型每次可以走不同的分支、跑不同的循环次数。

minigraph.py 里，VJP 规则只写了一遍（用 apply 写成），两种模式共用：在图模式下 apply 只加节点，在即时模式下 apply 直接计算。代码 01 的演示 1 用 TensorFlow 论文图 1 的那个 MNIST 分类器形状（784 → 100 → 10，批大小 32）对照了三者：

两种模式的梯度逐位相同，30 步训练每一步的损失也逐位相同（2.6130 → 0.0209）。结论很简单也很重要：自动微分的数学是一样的，两种系统的区别只在”什么时候执行”。

几个值得注意的细节：

• 图模式的前向图有 12 个节点，加上梯度子图后变成 30 个——反向本身也是图。系统可以像优化前向一样优化、切分、调度反向（例如把梯度的通信和反向计算重叠起来，第 12 讲）。
• 图模式每步执行 20 个算子，即时模式每步执行 22 个。多出来的 2 个是：我们的磁带实现对不需要梯度的输入 $x$ 也算了梯度（一次转置和一次矩阵乘法）。图模式的执行器只执行 fetches 依赖的节点，这个无用的梯度自然被剪掉了；真 PyTorch 则会检查每个输入的 requires_grad 来跳过它。这是”看得见全局”带来好处的一个小例子。
• PyTorch 论文专门提到了一个即时模式特有的难题：原地修改（mutation）。磁带记录的是”用哪个张量算出来的”，如果这个张量后来被原地改了，反向时读到的就是错的值。PyTorch 用张量版本号检测这种情况：大多数修改是安全的，危险的情况直接报错，提示用户改写程序，而不是悄悄地拷贝一份（那会引入难以察觉的性能悬崖）。

4. 计算图级优化 · Graph-level optimizations

4.1 两条代数等式

slides 用两条等式说明计算图级优化（这里统一用 $(n, o, h, w)$ 表示批、输出通道、空间位置）：

(1) Conv + BatchNorm 折叠。 卷积 $Z = \mathrm{Conv}(X, W) + B$，后面接推理模式的 BatchNorm $Y_{o} = (Z_o - \mu_o)/\sqrt{\sigma_o^2 + \epsilon} \cdot \gamma_o + \beta_o$。令每个输出通道的缩放 $s_o = \gamma_o / \sqrt{\sigma_o^2+\epsilon}$，则

也就是：把每个输出通道的卷积核乘以 $s_o$，再换一个偏置 $B'$，BN 就消失了。 $W \cdot s$ 和 $B'$ 只依赖常量，在编译时就能算好。省掉的是一个完整的逐元素内核——它要把整个激活张量读一遍、写一遍。

(2) 两个卷积合并。 卷积对卷积核是线性的，所以同一个输入上的两个卷积相加：

两个卷积 + 一个加法，变成一个卷积。FLOPs 减半，还省了一次激活的读写。

4.2 一个完整的例子：slides 里的 ResNet 小块（代码 02）

slides 第 5 页画了一个 ResNet 式的小块和它的”优化后的计算图”：

           ┌─ conv3x3 ─ BN ─┐
输入 x ────┤                ├─ add ─ relu            ⟹      x ─ conv3x3' ─ relu
           └─ conv1x1 ─ BN ─┘

代码 02 在 minigraph 上把这个改写过程一步步走了一遍（小形状上核对数值，ResNet-50 第二阶段的形状 32×64×56×56、BF16 上估算 H100 的代价）：

五步之后，6 个内核变成 2 个（一个 3×3 卷积 + relu），估算时间 73.9 → 23.4 µs（3.16×），每一步与原图的输出差都在 $10^{-15}$ 量级——这些改写是数学上等价的，误差只来自浮点舍入。

这张表里藏着两个比结果本身更重要的教训：

教训一：规则有前提，不满足就用不了。 第 ② 步合并失败了：两个卷积的核大小不同（3×3 和 1×1），规则 (2) 不适用。

教训二：好的优化有时要先经过一个”更差”的状态。 第 ③ 步把 1×1 卷积放大成 3×3（卷积核放在正中、四周填 0，padding 从 0 改成 1，输出完全不变）。单看这一步，FLOPs 从 8.23 G 升到 14.81 G（+80%）——一个以 FLOPs 为代价的优化器绝不会主动做它。（有意思的是，按 roofline 估算的时间几乎没变，50.5 → 50.6 µs：这个 1×1 卷积本来就是访存受限的，多出来的运算藏在了访存时间里；在计算受限的场景下，它就是实打实的变差。）但正是这一步让第 ④ 步的合并成为可能。只接受”代价立刻下降”的贪心优化器，会停在第 ② 步。 这就是第 10 讲 TASO 的出发点：它用代价驱动的回溯搜索，允许在搜索过程中暂时变差。

4.3 前提条件被违反会怎样？

代码 02 的演示 3 故意违反了两条规则的前提：

• 训练模式的 BN 不能折叠。 训练时，BN 用的是当前这一批数据自己的均值和方差，它们随输入变化，不是常量。用推理时的运行统计量硬套折叠，输出差了 22.70。所以 fold_conv_bn 要检查 BN 的四个统计量都是常量（推理模式、参数已冻结）。
• 中间有非线性就不能合并。 $\mathrm{relu}(\mathrm{Conv}(X, W_1)) + \mathrm{relu}(\mathrm{Conv}(X, W_2))$ 和 $\mathrm{relu}(\mathrm{Conv}(X, W_1 + W_2))$ 差了 12.76；而没有 relu 时差 $1.1 \times 10^{-14}$（线性，成立）。

4.4 为什么要”自动发现”图优化？

slides 指出，这些图改写规则传统上是专家手写的。问题在于：

1. 规则太多：算子种类、形状、硬件的组合爆炸，手写的规则永远不全；
2. 容易写错：每条规则都有前提条件（§4.3），漏检一个前提就是一个隐蔽的数值 bug；
3. 组合起来才有效：单条规则可能让代价变差（§4.2 的第 ③ 步），要搜索规则的组合。

所以第 10–11 讲的方向是：自动生成候选改写（图优化生成器）→ 自动验证它们是否等价（图优化验证器）→ 在这些已验证的改写上搜索最优的图（图优化器）。slides 的示意图正是这条流水线：从”ML 优化的数学性质”出发，生成候选优化，验证后得到可用的优化，再交给优化器使用。

5. 并行化概览 · Parallelization overview

slides 用 SGD 的更新公式引出并行化：$w \leftarrow w - \frac{\eta}{n}\sum_i \nabla \ell_i(w)$ 里有两个可以切的维度——数据（求和里的 $i$）和模型（$w$ 本身）。

5.1 数据并行：切数据

三步：① 把数据集切成若干批，每块 GPU 拿一份；② 每块 GPU 用完整的模型对自己那份数据做前向和反向；③ 把所有 GPU 的梯度聚合起来（AllReduce），每块 GPU 用同样的平均梯度更新参数。

它成立的理由就是 §2.2 的那个事实：损失是对样本取平均的，所以梯度也是。代码 04 演示 1：32 个样本平均切给 4 块”GPU”，各卡梯度取平均后和整批梯度差 $5.6 \times 10^{-17}$。一个容易踩的坑：如果各卡样本数不同（例如数据集最后一批不满，各卡分到 10、10、10、2 个样本），简单平均会错（差 0.17），要按样本数加权。

数据并行的代价有两个：每一步都要同步一份和模型一样大的梯度（第 12 讲的 AllReduce 算法），以及每块卡都存一整份模型和优化器状态（第 13 讲 ZeRO 要消除的冗余）。

5.2 张量并行：切矩阵

slides 用 Transformer 说明张量模型并行：把一个矩阵乘法 $Y = XW$ 切给多块 GPU。两种基本切法（代码 04 演示 2 都验证了）：

• 按列切 $W$：每块卡算 $Y$ 的一部分列，最后拼接（AllGather）；
• 按行切 $W$（同时按列切 $X$）：每块卡算一个形状完整的部分和，最后相加（AllReduce）。

Megatron-LM 的经典切法把两种配合起来：MLP 的第一个矩阵按列切，输出按列分布在各卡上，中间的激活函数各卡独立算；第二个矩阵按行切，正好接住这种分布，最后只需一次 AllReduce。代码验证它和不切的结果差 $2.2 \times 10^{-15}$；而如果第一个矩阵按行切、又在相加之前各自做 relu，结果差 1.20——非线性函数之前必须拿到完整的值，切法要和前后的算子配合。

5.3 流水线并行：切层

slides 的另一种模型并行：把前几层放在 GPU 1、后几层放在 GPU 2，中间结果在 GPU 之间传递。问题是任何时刻只有一块卡在干活。解决办法是把一批数据切成 $M$ 个微批次，让它们像流水线一样流过各个阶段。代码 04 演示 3 画出了 slides 上那张 2 卡 4 微批次的时间线：

GPU1 | F1 F2 F3 F4  .  . B4 B3 B2 B1
GPU2 |  . F1 F2 F3 F4 B4 B3 B2 B1  .

空闲的格子叫空泡（bubble），比例是 $(P-1)/(M+P-1)$：$P=2$、$M=1$ 时 50%（退化成朴素模型并行），$M=4$ 时 20%，$M=64$ 时 1.5%。

5.4 组合：3D 并行

slides 最后指出，真正训练大模型要把几种并行组合起来：Megatron-LM 组合数据并行和模型并行扩展到 512 块 GPU；DeepSpeed 的 “3D 并行”同时使用数据、张量、流水线三种。原则和第 1 讲 §3.1 一样——通信最频繁的切分放在最近的地方：张量并行放在一个节点内（NVLink），流水线并行跨节点，数据并行在最外层。怎样组合、各用多少，是第 15 讲自动并行要回答的问题。

6. ML 编译概览 · ML compilation overview

6.1 张量函数与”等价实现”

slides 引入了 ML 编译的两个基本概念：

• 张量：存放模型输入、输出和中间结果的多维数组；
• 张量函数：描述张量之间计算的函数。注意一个张量函数可以包含多个算子——例如”linear + relu”可以看作一个张量函数。

ML 编译要解决的问题是：同一个张量函数有非常多种等价的实现方式（分块大小、循环顺序、哪几个算子合在一起、数据怎么排布），要从中找出最好的。slides 列出的三个共同目标：最少的内存使用、最短的执行时间、最高的硬件利用率。第 1 讲 §8 的实验已经让你看到：同一个前馈块，逐算子执行 1.78 ms，理想下界 1.46 ms；同一个 relu(x*2+1)*0.5，融合前后差 4 倍。

6.2 厂商算子库及其局限

传统的做法是硬件厂商提供手工调优的算子库：NVIDIA 的 cuDNN（卷积等）、cuBLAS（矩阵乘法，例如 cublasSgemm()）、cuRAND、cuSPARSE。框架遇到卷积就调 cudnnConvolutionForward()，遇到矩阵乘法就调 cuBLAS。TensorFlow 和 PyTorch 论文都承认，它们的单机性能主要来自这些库：TensorFlow 和 Torch 用同一版 cuDNN，所以性能相差不到 6%；PyTorch 和其他框架用同一版 cuDNN / cuBLAS，所以相差不到 17%。

slides 指出了这种做法的两个问题：

1. 新算子无法立刻得到支持：研究者发明了一种新的注意力变体，在厂商库支持它之前，只能用一串已有算子拼出来，速度很慢；
2. 硬件越来越复杂，手写的内核越来越不是最优的：第 1 讲 §4.1 看到，每一代 GPU 都有需要专门利用的新特性，手写内核追不上。

6.3 自动代码生成

另一条路是让编译器自动搜索高性能的程序：TVM、Triton、Mojo、TileLang 等。slides 列出的好处：新算子立即可用；性能可能超过手写内核（因为搜索能覆盖人想不到的组合，也能针对具体形状和硬件定制）。第 1 讲 §8.5 里 torch.compile 为 relu(x*2+1)*0.5 自动生成的那个 C++ 循环，就是一个最小的例子。第 8 讲讲”按块编程”的 Triton 和 TVM，第 9 讲讲怎样自动调优。

7. 内存优化 · Memory optimizations

7.1 显存是训练的瓶颈

slides 回顾了一个事实：能训练的最大模型受限于 GPU 显存，而更大的模型往往效果更好。训练时显存里除了参数、梯度、优化器状态，还有激活：§3.2 说过，反向传播要用前向的中间结果，所以前向算出的所有激活都要活到反向用完为止。对深的网络，激活往往是显存的大头。片上存储越小（第 1 讲 §3.2），这个问题越尖锐。

7.2 重计算：用计算换显存（代码 03）

slides 的办法：前向时只保存部分节点（检查点），反向需要某一段的激活时，从最近的检查点把这一段重新算一遍（rematerialize）。如果 $N$ 层的网络每 $K$ 层存一个检查点：

对 $K$ 求最小，$K = \sqrt N$，显存 $O(\sqrt N)$。

代码 03 在一个 $N = 64$ 层的链式网络（每层 $h' = \tanh(hW)$）上实际数了一遍：

实测与公式完全一致：$K = 8$ 时只需存 16 个激活，省了 75%。代价是多做了 56 次层前向，约等于”再做一遍前向”——一次训练步大约是 1 份前向 + 2 份反向的计算，所以多出来的约是总计算的 1/3。再往两头看：$K$ 太大又回到 $O(N)$（整段都要暂存）；极端的”只存输入、每层反向都从头重算”只要 3 个激活，却要做 2144 次层前向（$O(N^2)$）。计算与显存之间是一条连续的权衡曲线。

两个容易被忽视的事实：

• 重计算是精确的：每种策略的梯度都和全存逐位相同（同样的运算、同样的顺序），不是近似。代码还和 PyTorch 自带的 checkpoint_sequential 对照，差 $4.5 \times 10^{-13}$（NumPy 和 PyTorch 的数学库在最后几位上的差别）。
• 真实网络不是一条链：有分支、有跳连，不同层的激活大小和重算代价也不同。怎样在一般计算图上选检查点，是一个优化问题（附加讲 A1）。实际系统里常用”选择性重计算”：只重算那些”占显存多、重算便宜”的激活，例如注意力里的 softmax 输出。

7.3 零冗余：用通信换显存

slides 指出，分布式训练里还有另一种浪费：冗余。数据并行的每块卡都存了一整份参数、梯度和优化器状态；模型 / 流水线并行也会复制一些中间张量。ZeRO 的核心想法是：把这些被复制的东西切开，每块卡只存 $1/P$，需要时再从别的卡广播过来——代价是额外的通信。第 13 讲会细讲它的三个阶段。

7.4 计算、内存、通信的三角

slides 用一个三角形总结：重计算在”计算”和”内存”之间权衡，零冗余在”通信”和”内存”之间权衡。几乎所有的训练系统优化，都是在这个三角里移动——没有免费的午餐，只有”把瓶颈换到不那么紧张的资源上”。

8. TensorFlow vs PyTorch：延迟执行 vs 即时执行 · Deferred vs eager

这是本讲的第二部分，也是两篇必读论文的主题。

8.1 TensorFlow 的延迟计算模型

TensorFlow 1.x 的程序分两个阶段（TensorFlow 论文图 1）：

# 第一阶段：构图（这里没有任何数值计算）
x = tf.placeholder(tf.float32, [BATCH, 784])
W_1 = tf.Variable(tf.random_uniform([784, 100]))
layer_1 = tf.nn.relu(tf.matmul(x, W_1) + b_1)
...
train_op = tf.train.AdagradOptimizer(0.01).minimize(loss)
# 第二阶段：反复执行这张图
with tf.Session() as sess:
    for step in range(NUM_STEPS):
        sess.run(train_op, {x: x_data, y: y_data})

slides 总结了这种模型的好处：

1. 图级优化：例如把 MatMul + Add + ReLU 合成一个 “Linear + ReLU”；
2. 延迟（惰性）执行：只有被请求的输出才会被计算——图里的”辅助层”如果这次没被用到，就不执行；
3. 利用全局信息做优化：因为执行前就能看到整个程序。

代码 01 的演示 2 在 minigraph 上把这三点都做了出来：

• 部分执行：一张带辅助输出（监控用的统计量）的图，只取损失时执行 6 个算子，同时取辅助输出时执行 9 个。
• 公共子表达式消除：用户在辅助分支里又写了一遍 x @ W1，CSE 发现它和主分支的那一次完全相同，去掉了一个重复节点（9 → 8 个算子）。
• 融合：批大小 8192 时，6 个内核（matmul、add、relu、matmul、add、softmax_xent）融合相邻逐元素算子后变成 4 个，H100 估算 31.1 → 22.0 µs。

这三件事有一个共同点：都需要在执行之前知道后面会发生什么。即时执行到 x @ W1 的那一刻，系统不知道后面还会再算一次，也不知道 add 和 relu 紧跟在后面。

slides 同样列出了代价：难调试（构图时打印出来的只是符号）；先构图再执行的写法不自然。代码 01 演示 3：图模式里打印 relu 的结果，得到的是一个符号 Node(relu_1: relu[4, 3])，要等 run() 才有数值；把符号张量用在 Python 的 if 里会直接报错——符号没有数值，Python 不知道该走哪个分支。

控制流怎么办？ 依赖数据的条件和循环必须变成图里的特殊算子。TensorFlow 论文借用了经典数据流机器里的 Switch / Merge 原语：Switch 根据条件把数据送到两个输出之一，没被选中的那一路收到一个特殊的”死”值，死值一路传播直到 Merge，Merge 把活着的那一路输出——这样只有被选中的分支真正执行。minigraph 的简化版 where(cond, a, b) 则两个分支都算（演示 3 里执行了 5 个算子，包括没用到的那个分支），可见”把控制流放进图里”本身就是一个需要专门设计的系统问题。

8.2 PyTorch 的命令式、Pythonic 模型

PyTorch 的模型就是一个普通的 Python 程序：每一行立刻执行，print、调试器、if、for、递归都能直接用。slides 的评价：

• 为生产力而不是性能优化；
• 容易原型开发和调试；
• 但因为没有全局信息，会错过一些优化——slides 的例子是：执行 matmul 的那一刻，relu 还没有出现，所以不能直接把 relu 融进 matmul。

PyTorch 论文的回答是：通过精心的工程实现，即时执行可以基本不牺牲性能——关键手段见 §9.2：C++ 核心、控制流和数据流分离（Python 只负责发射内核，GPU 异步执行）、缓存显存分配器、引用计数即时释放内存。

8.3 对比

slides 的对比表：

8.4 PyTorch 2.0：兼得两者？

slides 最后介绍了 PyTorch 2.0 的 torch.compile(model)：在保留即时执行编程模型的前提下，抓取计算图，再把它编译成优化的内核；slides 引用的结果是它在训练和推理上都带来了显著的加速（相对 eager 模式的加速比分布图）。

它的关键组件 TorchDynamo 在 Python 字节码层面拦截函数的执行，把能识别的张量运算记录成一张 FX 图，交给后端编译器（默认是 Inductor，第 1 讲 §8.5 生成 C++ 内核的就是它）。代码 01 演示 4 写了一个”只记录、不优化”的后端，看 Dynamo 交给后端的是什么：

• relu(x @ w1) @ w2：抓到 1 张图，里面是 [matmul, relu, matmul]；
• 含依赖数据的 if h.sum() > 0: 的函数：抓到 2 张图，发生 1 次 graph break，原因是 generic_jump TensorVariable()——条件跳转依赖一个张量的值，Dynamo 只能在这里把图切断，把这个判断交回 Python 执行；第一张图是 [matmul, relu, sum, gt]，第二张图是被选中的分支 [mul]。

这正好对应 §8.1 的控制流问题：TensorFlow 1.x 要求用户把控制流写进图里；PyTorch 2 让用户随便写，遇到写不进图的地方就”断图”，代价是图被切得越碎，编译器能做的全局优化越少。第 7 讲的 mega-kernel 和第 8 讲的编译器还会回到”图有多大”这个问题。

9. 论文精读 · Paper deep dive

9.1 TensorFlow: A System for Large-Scale Machine Learning（OSDI’16）

要解决什么问题？ Google 的上一代系统 DistBelief 采用参数服务器架构：无状态的 worker 做计算，有状态的参数服务器保存参数。它在三个方面不够灵活：① 定义新的层要用 C++ 实现，对研究者门槛太高；② 想改优化算法（例如新的更新规则、需要原子更新一组参数）就要改参数服务器的实现；③ worker 的执行模式是固定的”读数据和参数 → 前向 → 反向 → 写回梯度”，装不下循环网络、对抗网络、强化学习这类训练流程。此外它只为大集群设计，很难缩小到单机工作站或手机上。

为什么难？ 要同时满足：研究者的灵活性（新层、新优化器、新训练算法）、生产的性能和规模（数百台带 GPU 的服务器）、部署的多样性（从数据中心到手机，再到 TPU 这样的专用加速器）。

核心想法： 用一张统一的数据流图同时表示计算、共享状态和修改状态的操作。

• 原始算子级的数据流图：图的节点是矩阵乘法、卷积这样的单个数学算子，而不是 DistBelief 那样粒度很粗的”层”，所以用户可以在脚本里自由组合新层，系统也容易对它们自动求导。
• 可变状态是图的一部分：变量（参数）和更新变量的操作（例如 AssignAdd）都是图里的节点。于是优化算法不再是写死在参数服务器里的代码，而是用户可以自己写的子图。论文的一个推论很有意思：TensorFlow 里没有所谓的”参数服务器”——所谓 PS 任务只是”放着变量的那些任务”，它们可以运行任意的子图（例如把一部分计算挪到参数所在的地方，减少网络流量）。
• 延迟执行：先定义整个程序，再执行它的优化版本；执行时可以利用图的依赖结构一次性向 GPU 发射一串内核，不必等中间结果。代价是更复杂的特性（例如动态控制流）也要放进图里——于是有了 Switch / Merge。
• 异构设备的统一抽象：一个设备只需实现”发射内核、分配内存、和主机之间拷贝缓冲区”三个方法；每个算子可以为不同设备提供不同实现，同一个程序可以在 GPU、TPU、手机 CPU 上运行。
• 部分与并发执行：每次 run 只执行请求的子图，多个子图可以并发执行在同一张图上（例如训练、检查点、数据预处理同时进行）。
• 分布式执行：图被切分到各个设备上，跨设备的边自动替换成发送 / 接收（Send / Recv）节点；放置算法根据用户的部分约束和设备能力决定每个算子放在哪里。

扩展性的几个案例（都是”把它写成子图”的直接结果）：符号微分（梯度也是子图，用户可以自定义某些算子的梯度）；超大的嵌入矩阵切分到多个 PS 任务上，用 Gather / Part / Stitch 实现稀疏读取和更新；用户级的检查点（Save / Restore 也是图里的算子，用户可以定制保存策略）；同步 / 异步训练的协调，以及缓解掉队者（straggler）的备份 worker——一步训练只等前 $m$ 个 worker 的梯度，丢弃最慢的几个。

结果好不好？

• 单机（一块 Titan X）上训练四个卷积网络：比 Caffe 快，和最新的 Torch 相差在 6% 以内（两者用同一版 cuDNN）；Neon 在三个模型上更快，因为它用了汇编手写的卷积内核。
• 同步训练的协调开销：一个”空”训练步，只取一个 4 字节标量时，中位数步时间从 1 个 worker 的 1.8 ms 增长到 100 个 worker 的 8.8 ms；取 100 MB 的稠密模型时从 147 ms 增长到 613 ms。
• Inception-v3 图像分类：和同期的参数服务器系统 MXNet 相比略快（两者用同一版 cuDNN）；增加到 200 个 worker 时吞吐增长到约 2300 张图/秒，但收益递减；同步训练的中位数步时间比异步长约 10%，但尾部（第 90 百分位以上）被掉队者拖得很严重。在 50 个 worker 上加备份 worker：加 4 个时步时间最短（1.93 s），按资源归一化后加 3 个最划算（约 9.5% 的加速）。
• 语言模型（One Billion Word，LSTM-512-512，词表 4 万）：增加 PS 任务数能提高吞吐（把 softmax 的计算分散到 PS 上）；采样 softmax 减少了传输和计算，吞吐进一步提高。

局限与讨论： 论文自己承认：静态数据流图对深度强化学习这类算法开始显得束缚；默认的优化策略还没有找到对所有用户都好的设置，作者正在研究自动放置、自动融合、内存管理和调度。回过头看，这几条正好是后来十年 ML 系统研究的主线（第 8–15 讲）。从读者的角度还可以追问：评估主要是”系统吞吐”，没有报告”达到同样精度要多久”（time-to-accuracy），作者也明确把它留给了别的论文。

9.2 PyTorch: An Imperative Style, High-Performance Deep Learning Library（NeurIPS’19）

要解决什么问题？ 当时的主流框架（Caffe、CNTK、TensorFlow、Theano）都构建静态数据流图：执行前能看到整个计算，理论上有利于性能和扩展，但代价是易用性、可调试性和可表达的计算类型。已有的即时执行框架要么牺牲性能（Chainer），要么用了更快但表达力更弱的语言（Torch 用 Lua、DyNet）。论文要证明：易用性和速度并不矛盾。

四条设计原则（转述）：

1. Pythonic：做 Python 生态的一等公民，接口简单一致，和标准的画图、调试、数据处理工具自然配合；
2. 研究者优先：把机器学习固有的复杂性藏在库的内部，接口直观、没有副作用、没有意外的性能悬崖；
3. 务实的性能：性能要有竞争力，但不能以易用性为代价——论文的说法大意是：为了明显更好用而牺牲 10% 的速度可以接受，牺牲一倍不行；同时给研究者手动控制执行的工具；
4. “更差就是更好”（worse is better）：工程资源固定时，保持内部实现简单，省下的时间可以用来加功能、跟上领域的变化。

易用性的设计：

• 模型就是 Python 程序：层是”带隐式参数的有状态函数”，写成 Python 类；GAN 这种两个模型、两个损失交替训练的流程，在刚性的 API 里很别扭，在 PyTorch 里就是几行普通代码；
• 互操作：torch.from_numpy() 和 .numpy()、DLPack 格式交换数据时零拷贝——两边只是对同一块内存的不同解释；
• 可扩展：自定义可微函数（继承 torch.autograd.Function，实现 forward 和 backward）、自定义数据集（实现 __getitem__ 和 __len__），任何组件都可以替换；
• 自动微分：Python 是动态语言，事前的源到源变换很麻烦，所以 PyTorch 用运算符重载在每次执行时记录计算（磁带），做反向模式 AD；并用张量版本号支持对原地修改的求导（见 §3.3）。

性能的设计（核心挑战：Python 解释器有全局锁 GIL，开销大）：

• 高效的 C++ 核心（libtorch）：张量、算子、自动微分的求值器都在 C++ 里，反向传播在多线程求值器里执行，不需要持有 GIL；
• 控制流与数据流分离：控制流（分支、循环）由 Python 和 CPU 上的 C++ 处理，产生一串算子调用；算子通过 CUDA stream 异步地排进 GPU 的队列。只要每个算子在 GPU 上跑得够久，CPU 就能跑在 GPU 前面，GPU 几乎一直是满的；
• 自定义缓存分配器：GPU 上的 cudaFree 可能阻塞调用者直到之前排队的所有工作完成；PyTorch 自己缓存显存、重复分配，并按 512 字节取整减少碎片，每个 CUDA stream 一个内存池（利用”同一个 stream 上的操作是串行的”这一点保证安全）；
• 多进程：扩展 Python 的 multiprocessing，把发给其他进程的张量自动放进共享内存，而不是序列化；
• 引用计数：即时执行的库不知道一个张量将来还会不会被用到。垃圾回收会推迟释放、增大显存占用（论文举例说 Torch7 的用户常常在程序里到处手动触发 Lua 的垃圾回收）；PyTorch 用引用计数，计数归零立刻释放。

结果好不好？

• 异步数据流：ResNet-50 的执行时间线显示，CPU 排队工作的速度远快于 GPU 执行，GPU 执行时间大约是 CPU 调度时间的 3 倍，设备利用率接近满载；
• 内存管理：第一次迭代被 cudaMalloc / cudaFree 拖慢，之后缓存分配器开始复用，这个开销消失；
• 基准：在 AlexNet、VGG-19、ResNet-50、MobileNet、GNMTv2、NCF 六个模型上，和 CNTK、MXNet、TensorFlow、Chainer、PaddlePaddle 比较，PyTorch 都在最快框架的 17% 以内——作者把它归因于大家调用的是同一版 cuDNN / cuBLAS；
• 采用度：arXiv 上提到 PyTorch 的论文比例持续上升。

局限与讨论： 论文的性能论证有一个前提：每个算子足够大，Python 的开销才能被藏在 GPU 执行时间后面。当算子很小（例如 LLM 解码时 batch=1 的逐元素运算），CPU 就跟不上 GPU 了——第 1 讲 §4.4 测到的”每次调用 0.3 µs 的固定开销”和 §8 的”解码时 31% 是启动开销”正是这种情况。论文的”未来工作”一节提到了 PyTorch JIT（让程序脱离 Python 解释器执行、做进一步优化）——后来演化成 TorchScript，再到 PyTorch 2 的 torch.compile。

9.3 两篇论文的联系

1. 同一个问题的两个端点。 两篇论文都在回答”怎样让研究者方便地写模型、又让系统高效地执行”，但押注的方向相反：TensorFlow 押注”系统看见整个程序”，PyTorch 押注”程序员看见整个程序”。本讲 §8 的实验说明，两者的自动微分数学完全一样，差别全在执行时机。
2. 各自的短板正是对方的长处。 TensorFlow 论文结尾承认静态图束缚了动态算法；PyTorch 论文的性能依赖”算子足够大”和厂商库。前者的部分执行、CSE、融合、自动放置，后者需要 JIT / 编译器才能拿回来。
3. 后来的趋同。 TensorFlow 2 默认即时执行、用 tf.function 把需要加速的部分转成图；PyTorch 2 用 Dynamo 从即时代码里抓图。JAX 走了第三条路：用户写函数式的 NumPy 代码，jit 通过追踪把它变成图交给 XLA。三者殊途同归：用户写即时风格的代码，系统在背后找机会构图。
4. 一个可以写进评审的新想法。 TensorFlow 论文里”可变状态也在图里、优化器是用户写的子图”的思想，在 PyTorch 2 时代有了新的意义：如果编译器能把优化器的更新也抓进图里，就能把”反向的最后一个算子”和”参数更新”融合起来，甚至在梯度一算出来就开始通信和更新（第 12–13 讲会看到这类优化）。这是”用一篇论文的想法解决另一篇论文的问题”的一个例子。

10. 小结 · Summary

1. 深度学习系统的标准部件：张量与算子、计算图、自动微分、图级优化、并行化、数据布局与放置、内核优化（ML 编译）、内存优化。
2. 自动微分 = 每个算子的 VJP + 按逆拓扑序遍历。 符号微分（图）和磁带（即时）用的是同一套数学；minigraph 里两种模式的梯度逐位相同，和 PyTorch 差 $10^{-17}$ 量级。
3. 图级优化靠代数等式：Conv+BN 折叠、卷积合并等。slides 的例子经过五步改写从 6 个内核变成 2 个（H100 估算 3.16×）。规则有前提条件；好的优化可能要先经过一个更差的状态——所以需要搜索和验证（第 10–11 讲）。
4. 并行化：数据并行依赖”梯度是样本梯度的平均”；张量并行的切法要和非线性算子配合；流水线的空泡比例是 $(P-1)/(M+P-1)$。
5. ML 编译：同一个张量函数有大量等价实现；厂商库对新算子和新硬件反应慢，自动代码生成是出路。
6. 内存优化：重计算用计算换显存，$K=\sqrt N$ 时显存 $O(\sqrt N)$、多约一次前向，并且是精确的；零冗余用通信换显存。
7. 延迟执行 vs 即时执行：前者能剪枝、去重、融合、放置，但难调试、控制流要特殊处理；后者好写好调，但看不到全局。PyTorch 2 的 torch.compile 从即时代码里抓图，遇到依赖数据的控制流就断图。

11. 常见误区 · Common misconceptions

1. “PyTorch 没有计算图。” 有，只是它是在执行时动态记录的（磁带），每次前向都重新记录一遍；反向传播就是沿着这张图走。区别在于这张图在执行之后才完整，系统来不及用它做前向的优化。
2. “符号微分和自动微分是两回事 / 符号微分会表达式爆炸。” 数学课上的”符号微分”是对公式推导导数表达式；TensorFlow 的”符号微分”是在计算图上按 VJP 规则构造梯度子图，本质上仍是反向模式自动微分，中间结果共享，不会爆炸。
3. “图优化只是锦上添花。” 对访存受限的算子（BN、逐元素运算），融合 / 折叠省掉的是整个内核的读写，§4.2 的例子里能带来 3 倍的估算加速。
4. “Conv+BN 折叠任何时候都可以做。” 只在推理模式成立（§4.3）。
5. “数据并行把梯度加起来就行。” 要和损失的定义一致：损失是平均的，就要平均；各卡样本数不同，就要加权（§5.1）。
6. “重计算是一种近似。” 不是，重算的激活和原来逐位相同（§7.2）；它只是多花计算。
7. “即时执行一定比图执行慢。” 只要算子够大，异步执行能把 Python 开销完全藏起来（PyTorch 论文的结论）；只有在小算子、开销受限的场景下才明显吃亏。反过来，图执行如果没有做任何优化，也不会更快。
8. “torch.compile 能把任何 PyTorch 代码变成一张图。” 遇到依赖数据的控制流、不支持的 Python 特性、某些副作用，Dynamo 会断图，把这部分交回 Python 执行（§8.4）。

12. 练习 · Exercises（含参考答案）

练习 1（VJP）：写出 $y = \mathrm{sigmoid}(x) = 1/(1+e^{-x})$ 的 VJP，并说明反向时需要保存前向的什么。

答：$\mathrm{sigmoid}'(x) = y(1-y)$，所以 $\bar x = \bar y \odot y \odot (1 - y)$。只需保存输出 $y$，不需要保存 $x$。（这也是一个内存优化的小技巧：选择用输出表示导数的写法，可以早点释放输入。）

练习 2（广播的反向）：$Y = X + b$，$X$ 形状 $(B, d)$，$b$ 形状 $(d,)$。$\bar b$ 是什么？在 minigraph 里是哪个函数负责的？

答：$b$ 被广播了 $B$ 次，每一份都对损失有贡献，所以 $\bar b = \sum_{i=1}^{B} \bar Y_{i,:}$，即沿第 0 维求和。minigraph 里 _unbroadcast 调用 sum_to 算子完成这件事。

练习 3（Conv+BN 折叠）：某输出通道 $\gamma = 2$，$\beta = 0.5$，$\mu = 1$，$\sigma^2 = 3$，$\epsilon = 1$，卷积原来的偏置 $B = 0$。写出折叠后的缩放和新偏置。

答：$s = 2/\sqrt{3 + 1} = 1$，卷积核乘以 1（不变）；$B' = (0 - 1) \times 1 + 0.5 = -0.5$。

练习 4（卷积合并的代价）：§4.2 里第 ③ 步把 1×1 卷积放大成 3×3，FLOPs 增加了多少？为什么最终结果反而更好？

答：1×1 卷积的 FLOPs 变成原来的 9 倍（$2 \cdot 32 \cdot 64 \cdot 56^2 \cdot 64 \approx 0.82$ G → 7.4 G），总 FLOPs 从 8.23 G 升到 14.81 G。但放大之后两个卷积形状一样，可以合并成一个（7.41 G），还省掉了 add 和一次激活的读写；最终 FLOPs 比原图还少，内核数从 4 降到 2。

练习 5（重计算）：一个 $N = 100$ 层的链，每层激活 1 GB。(a) 全存要多少激活显存？(b) 最优的 $K$ 和对应的显存？(c) 多做多少层前向？

答：(a) 约 101 GB（101 个激活，含输入）。(b) $K = 10$，峰值 $N/K + K = 20$ 个激活，约 20 GB（处理最后一段时：前面 9 个检查点 + 这一段的起点和终点 + 段内重算的 9 个）。(c) 每段重算 $K - 1 = 9$ 层，共 10 段，多做 90 次层前向（≈ 0.9 次完整前向；总共 190 次）。把 03_rematerialization.py 的 N_LAYERS 改成 100 运行，得到的正是 101 / 20 个激活、100 / 190 次前向。

练习 6（流水线空泡）：4 个流水线阶段，要把空泡比例压到 10% 以下，至少要几个微批次？

答：$(P-1)/(M+P-1) < 0.1$，即 $3/(M+3) < 0.1$，$M > 27$，至少 28 个微批次。

练习 7（延迟 vs 即时）：下面哪些优化必须有”执行前的全局视图”？(a) 把 matmul 后面的 relu 融进 matmul 内核；(b) 用缓存分配器复用显存；(c) 剪掉没被请求的输出分支；(d) 引用计数释放不再使用的张量。

答：(a)、(c) 需要。(b)、(d) 即时执行也能做（PyTorch 论文就是这么做的）。不过如果有全局视图，内存还能规划得更好（提前知道每个张量的生命周期，静态分配，避免碎片）。

练习 8（graph break）：把代码 01 演示 4 的 branchy_torch 改成用 torch.where(h.sum() > 0, h * 2, torch.exp(h))，再用 torch._dynamo.explain 看看还有没有 graph break。两种写法各有什么代价？

答：改成 torch.where 后没有 Python 层面的分支，Dynamo 能抓成一张图、没有 graph break；代价是两个分支都会被计算（和 minigraph 的 where 一样）。原写法只算一个分支，但被切成两张图，编译器的优化范围变小。在分支很便宜时 where 更好；在分支很贵时，原写法或者 torch.cond 这样的控制流算子更好。

13. 延伸阅读 · Further reading

• 两篇必读论文：TensorFlow（OSDI’16）、PyTorch（NeurIPS’19）。
• A. G. Baydin et al., Automatic Differentiation in Machine Learning: a Survey（JMLR 2018）：自动微分最好的综述。
• T. Chen et al., Training Deep Nets with Sublinear Memory Cost（2016）：$\sqrt N$ 检查点的原始论文。
• J. Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation（ASPLOS’24）：TorchDynamo 和 TorchInductor 的论文，slides 最后那张加速比分布图的出处。
• R. Frostig, M. Johnson, C. Leary, Compiling machine learning programs via high-level tracing（SysML 2018）：JAX 的设计。
• M. Shoeybi et al., Megatron-LM（2019）：§5.2 的切法，第 14 讲的必读论文。
• 本仓库 CMU11-868 第 5 讲：从零实现 miniTorch 的自动微分，和本讲 §3 互补。

14. 下一讲预告 · Next lecture

第 3 讲《GPU 架构与 CUDA 编程（一）》从 §6 的”内核”往下走一层：GPU 为什么是 SIMD / SIMT 的？CUDA 程序的 grid / block / thread 层次怎样映射到 SM 和 warp？分支发散为什么慢？共享内存为什么能减少访存？GPU 怎样把成千上万个线程块调度到上百个 SM 上？配套代码会写一个纯 Python 的 SIMT 模拟器 gpusim.py，把 slides 里的每个 CUDA 例子都”跑”一遍，并数出访存事务、发散的 warp 和调度的波次。