第 15 讲 · 高级 ML 并行化：自动并行 · Automated ML Parallelization

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 15 · Fri 10/24 · Advanced ML Parallelization: Automated ML Parallelization。 必读论文：Z. Jia, M. Zaharia, A. Aiken, Beyond Data and Model Parallelism for Deep Neural Networks（FlexFlow，MLSys 2019）；L. Zheng et al., Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning（OSDI 2022）。 依据说明：本讲以原课 slides 14-auto-parallelization.pdf 为主线：

• 三种并行的比较表与”训练大模型要组合它们”；
• 手工设计并行策略的问题（难以设计和实现、性能不优、可移植性差）；
• FlexFlow：“搜索空间 + 代价模型 + 搜索算法”；SOAP 四个维度（Sample、Operator、Attribute、Parameter）及其混合；
• 发现好策略的两个挑战（策略数以十亿计 → MCMC；在硬件上评估太慢 → 执行模拟器）；
• 模拟器的关键观察（算子性能高度可预测、模型只有少数几种算子）和精度（相对误差 30% 以内、保持排序）；DLRM 案例；
• Alpa：算子间与算子内并行、两级层次搜索空间、分层优化（算子间用动态规划、算子内用 ILP）、运行时编排、评估与宽 ResNet 案例；
• 讨论题：比较 FlexFlow 和 Alpa。

Alpa 部分 slides 主要是论文的图，本讲按论文正文补全了分片规格、重分片、ILP 目标和动态规划公式。论文数字只做转述。

flexflow_mini.py（SOAP 切法、任务图、执行模拟器、MCMC）、alpa_mini.py（分片规格、重分片代价、算子内最优切法、算子间动态规划）和全部实验都是本讲义自己设计的缩小版，不是两个系统的复现。 配套代码 Code：code/lecture15/（4 个脚本 + 两个共享模块；只需 numpy；最长的约 25 秒） 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标

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

1. 用 SOAP 的四个维度描述一个算子的切法，说出每个维度切开时要付什么代价（同步参数、搬输入、光晕）。
2. 解释 FlexFlow 的执行模拟器怎样把”算子图 + 拓扑 + 策略”变成任务图并模拟，它依赖哪些假设，为什么”保持排序”比”误差小”更重要。
3. 写出 Metropolis–Hastings 的接受准则，解释 β 的作用，以及 MCMC 为什么能跳出局部最优。
4. 读懂 Alpa 的分片规格（如 S0R、RS1、S01R），推出两种规格之间的重分片要什么通信。
5. 写出 Alpa 算子内的目标函数（节点代价 + 边上的重分片代价）和算子间的动态规划，解释两层为什么这样分。
6. 比较 FlexFlow 和 Alpa：搜索空间、代价模型、搜索算法、各自的局限。

0.2 先修与本讲范围

0.3 本讲地图

§1  为什么要自动并行
§2  FlexFlow 的 SOAP 搜索空间                         ← 代码 01
§3  执行模拟器：任务图、完整模拟、增量模拟、精度         ← 代码 02
§4  MCMC 搜索与结果                                    ← 代码 03
§5  Alpa（一）：两级层次空间与算子内并行                ← 代码 04（演示 1–3）
§6  Alpa（二）：算子间动态规划、运行时、评估             ← 代码 04（演示 4–5）
§7  讨论：比较 FlexFlow 和 Alpa
§8  论文精读：FlexFlow；Alpa
§9 小结  §10 常见误区  §11 练习  §12 延伸阅读  §13 下一讲预告

1. 为什么要自动并行 · Why automate

slides 先重复了第 14 讲的比较表：

• 数据并行可大规模并行，但放不下大模型、参数多时扩展差；
• 张量并行能训大模型，但扩展度有限、要传中间结果；
• 流水线并行对深的模型高效，但有空泡。

训练大模型要把它们组合起来。手工设计组合有三个问题（slides）：

• 难以设计和实现：每个模型、每种硬件都要专家重新想一遍，还要写分布式代码；
• 性能不优：人能想到的通常是”整层都用同一种并行”；
• 可移植性差：换一个模型或集群，策略就要重做。

第 14 讲 05 号脚本枚举了 GPT-3 的 30 种 (t, p, d) 组合，就已经需要一个代价模型来挑；如果每个算子都可以有自己的切法，组合数会大到无法手工处理。这一讲的两个系统把问题变成：定义搜索空间 + 能快速评估的代价模型 + 搜索算法——和第 9 讲（自动调优）、第 10 讲（图替换）是同一个套路，只是对象换成了并行策略。

2. FlexFlow 的 SOAP 搜索空间 · The SOAP search space

2.1 四个维度

FlexFlow 用输出张量怎样切分来描述一个算子怎样并行：每块设备计算输出的一个不相交的子集。

• Sample（样本）：沿 batch 维切——数据并行；
• Operator（算子）：不同的算子放到不同的设备上——模型并行（算子之间）；
• Attribute（属性）：沿样本内部、不涉及参数的维度切，例如图像的像素、图的顶点、序列的位置；
• Parameter（参数）：沿要切开参数的维度切，例如输出通道——张量并行。

一个算子的并行配置 c 给出每个可切维度的并行度（切成几份，均分），以及每个子任务放在哪块设备上；一个策略 S 为每个算子选一个配置，各算子之间独立选择。数据并行、模型并行、专家的混合策略都是 SOAP 空间里的点。

2.2 同一个计算，不同的代价（代码 01 演示 1）

slides 用一维卷积说明各个维度。代码把一个一维卷积（n = 8、c_in = 16、c_out = 32、L = 64、k = 3）按 6 种方式切到 4 块卡上，结果全部和单卡相同，但代价不同：

（整个输入 8,192 个元素，整个权重 1,536 个。）

• 样本维：每卡只读 1/4 的输入，但存完整的权重，迭代末尾要同步梯度；
• 参数维：每卡只存 1/4 的权重，不用同步，但要读完整的输入；
• 属性维：和样本维一样复制权重，还要多读 k − 1 个位置的”光晕”（相邻块边界上的数据）。

slides 还展示了一个”比数据并行快 1.2 倍”的 SOAP 混合策略：不同的层用不同的维度组合。

2.3 空间有多大（代码 01 演示 2–3）

代码把每个维度的并行度限制为 2 的幂、设备按块对齐，8 块卡上：

• AlexNet（8 个算子）有 3.3 × 10¹¹ 种策略；
• 一个 DLRM 式模型（16 个算子）有 3.7 × 10²² 种。

真实的 SOAP 空间允许任意设备分配，更大。slides：“SOAP contains billions or more possible strategies”。

不同策略的代价差别巨大。AlexNet（batch 256、2 节点 × 4 卡）的模拟时间：数据并行 595.6 ms，模型并行 315.6 ms，专家策略（卷积数据并行、全连接切参数）72.0 ms，三个随机策略 334–754 ms。好策略稀少，需要有方向的搜索。

3. 执行模拟器 · The execution simulator

3.1 为什么能模拟

slides 的关键观察：

• ML 算子的性能高度可预测：稠密矩阵运算的时间几乎只取决于形状，不取决于数值；
• 模型只包含少数几种不同的算子：每种算子、每种输出大小只需要真跑一次，结果缓存起来。

FlexFlow 的四条假设：

• A1：任务的执行时间可预测、方差小、和输入内容无关；
• A2：带宽 b 的连接上传 s 字节耗时 s/b（能用满带宽）；
• A3：每个设备按 FIFO 处理分给它的任务（GPU 就是这样）；
• A4：运行时开销可以忽略。

3.2 任务图（代码 02 演示 1）

给定算子图、设备拓扑和策略，构建任务图：

1. 每个算子按它的配置拆成 |c| 个计算任务；
2. 对每条数据依赖，算出生产者的每个子任务写了输出的哪一块、消费者的每个子任务要读输入的哪一块；有重叠的两个任务如果在同一块设备上，直接连一条边，否则插入一个通信任务，放在两块设备之间的连接上。

FlexFlow 把每条硬件连接也当作一个”设备”，所以计算和通信可以自然地重叠。代码还加了参数同步任务（同一份参数分片的多个副本）。代码仿照论文图 5 的三层网络（嵌入 → 线性 → 线性）打印了全部 15 个任务的就绪、开始、结束时间。

3.3 完整模拟与增量模拟

完整模拟（论文算法 1）：任务就绪后进入一个按就绪时间排序的优先队列；出队时，开始时间 = max(就绪时间, 设备上一个任务的结束时间)，结束时间往后传给后继。这是 Dijkstra 最短路的一个变体，一次迭代的时间 = 最后一个任务的结束时间。

增量模拟（delta simulation，算法 2）：MCMC 每次只改一个算子的配置，时间线的大部分不变。增量模拟从上一张任务图出发，只更新被改动的任务，再像 Bellman–Ford 一样把变化往后传播。结果和完整模拟完全相同。论文测得它比完整模拟快 2.2–6.9 倍，卡越多加速越大。（代码没有实现它，见代码 README。）

3.4 准不准（代码 02 演示 2）

没有 GPU 集群可以真跑，代码用一个”更细的执行模型”当作”真实执行”，故意违反 A1、A2：

• 计算时间有随机波动；
• 每个任务多 20 µs 的启动开销；
• 小消息用不满带宽：效率 = 字节数 / (字节数 + 1 MB)。

对每个模型评估 300 个策略：

如实地读这张表：

• AlexNet 的通信都是大消息，“能用满带宽”基本成立，理想模拟器既准又保持排序——这正是论文报告的两件事（误差在 30% 以内，保持真实执行时间的排序）。
• DLRM 的整体排序也不错，但最好的那些策略之间排序几乎是乱的。原因：DLRM 的好策略由大量小消息组成（每张嵌入表的输出只有几百 KB），恰好是 A2 失效的地方。搜索会主动走向模拟器最乐观的区域，于是模型误差被放大。
• 按消息大小校准带宽后排序恢复。但这一行几乎完美是”必然”的：校准版和”真实执行”共用同一个确定性模型，只差随机波动。它说明的是”假设对了有多重要”，不是”怎样校准”。

教训：评估一个代价模型，不只看平均误差，还要看搜索会去的地方它准不准。（第 9 讲的 Ansor 也遇到过类似的问题：代价模型在”好程序”附近的相对排序最重要。）

3.5 快不快（代码 02 演示 3）

代码的 Python 模拟器评估一个策略约 0.2 ms（AlexNet）/ 0.6 ms（DLRM），而被模拟的迭代本身要几十到几百毫秒——快 160–2,600 倍，还没算真跑时部署策略的时间。论文的模拟器比真实执行快三个数量级。

4. MCMC 搜索与结果 · MCMC search

4.1 Metropolis–Hastings

把代价变成一个概率分布：p(S) ∝ exp(−β · cost(S))。从某个策略出发，每一步提议一个新策略 S*：随机选一个算子，把它的配置换成一个随机的配置（这个提议分布是对称的）。接受概率：

$\alpha(S \to S^*) = \min\big(1,\ \exp(\beta \cdot (\text{cost}(S) - \text{cost}(S^*)))\big)$

更好的提议总被接受；更差的以一定概率被接受——MCMC 倾向于像贪心一样往低代价走，但也能跳出局部最优。初始策略用数据并行、专家策略和随机策略；每个初始策略搜到时间预算用完，或者”一半时间没有改进”就停。（代码把代价差除以初始代价，让 β 和模型的时间尺度无关。）

4.2 AlexNet 与 DLRM（代码 03 演示 1–2）

找到的策略：

• AlexNet：卷积层基本按样本切，全连接层不再在 8 块卡上复制，而是集中在一个节点内的几块卡上、切参数——专家直觉的”加强版”；
• DLRM：嵌入表不复制，每张放在一块或少数几块卡上，只传查出来的向量。数据并行要对 2 GB 的嵌入表做跨节点 AllReduce，而每个样本只查每张表的一行，几乎所有同步都是浪费的。

slides 的 DLRM 图显示 FlexFlow 比数据并行快约 10 倍；这里的 770 倍极端得多（表更大、网卡更慢、数据并行用稠密 AllReduce 同步嵌入梯度），数字不可比。

4.3 好策略好在哪里（代码 03 演示 5）

把一次迭代拆开（各类任务时间之和）：

• 数据并行：激活传输 0，参数同步巨大（AlexNet 250 MB、DLRM 2 GB）；
• 专家策略和搜索结果用一些激活传输换掉了大部分参数同步——第 14 讲”传激活还是传参数”的取舍，现在逐个算子、自动地做。

搜索结果的总计算时间也更少（AlexNet 185 vs 195 ms，DLRM 13 vs 34 ms）：有些算子不再切成 8 份，任务更大、GPU 效率更高。论文在 NMT 上观察到同样的两件事：数据传输少 2–5.5 倍，总计算时间比数据并行少 20%。

4.4 能搜到最优吗？β 的作用（代码 03 演示 3–4）

小网络（3 个线性层）的空间只有 17,576 种策略，可以穷举出全局最优（1.302 ms）。MCMC 每次 1,500 步（不到空间的 1/10）、10 个种子：3 个找到了全局最优，其余都在最优的 2% 以内。论文在 4 块卡上的 LeNet、RNNLM 上做过同样的检查，报告找到了全局最优。问题本身是 NP 难的（可以从最小完工时间问题归约），MCMC 不保证最优。

β 的扫描（AlexNet，从数据并行出发，2000 步，5 个种子平均）：

β 太小等于乱走；这个问题上 β = 500 最好，贪心只差一点（局部最优不多）。

5. Alpa（一）：两级层次空间与算子内并行 · Alpa: intra-operator parallelism

5.1 两类并行

Alpa 不按”数据 / 张量 / 流水线”分类，而按是否切开算子分：

• 算子内并行（intra-operator）：把一个算子切到多块设备上——数据并行、张量并行、ZeRO 都是它的特例（只是切的维度不同）。设备之间要做集合通信（AllReduce、all-gather……），通信量大，适合高带宽的连接；
• 算子间并行（inter-operator）：不同的算子放到不同的设备上——流水线并行。设备之间只在级边界点对点传激活，通信少，但有空泡。

关键想法（slides：“two-level hierarchical search space”）：集群天然是分层的（节点内高带宽、节点间低带宽），并行也分两层——算子间并行映射到节点之间，算子内并行映射到节点内。这样把一个巨大的联合空间拆成两个可以分别（近似）最优求解的子问题。

5.2 设备网格与分片规格（代码 04 演示 1）

算子内并行在一个二维的逻辑设备网格上进行（例如 2 × 4：轴 0 两行、轴 1 四列）。一个张量的分片规格（sharding spec）说明它的每一维沿哪些网格轴切：

• S0：沿轴 0 切；S1：沿轴 1 切；S01：沿两个轴一起切；R：复制。
• 二维张量写成”行 + 列”，例如 S0R：行沿轴 0 切、列复制；RS1：列沿轴 1 切。

相邻算子要求的规格不同时，就要重分片（resharding）。代码按”逐个网格轴比较源和目标”的规则推导，和 Alpa 论文表 2 一致：

同样的操作，沿跨节点的轴比节点内慢一个数量级。

5.3 一个算子的切法（代码 04 演示 2）

Alpa 为 XLA 的每种基本算子手工枚举了并行算法（论文表 3 列出了矩阵乘法的几种）。思路是把网格的每个轴分给算子的一个循环维。对 Y[b, j] = Σᵢ X[b, i]·W[i, j]，每个轴可以分给：

• b：数据并行，反向要同步 W 的梯度；
• i：行并行，前向要 AllReduce Y；
• j：列并行，反向要 AllReduce dX。

两个轴共 9 种组合（重算子不允许复制计算）。一个 8192 token × 4096 → 4096 的 Linear 在 2 × 4 网格上：最便宜的是”慢的轴 0 做数据并行、快的轴 1 做张量并行”（0.55 ms），两个轴都做同一件事的最贵（2.5–4.8 ms）。但单个算子最便宜不等于整个图最便宜，还要算重分片。

5.4 ILP：节点代价 + 边代价

Alpa 把整个计算图的切法选择写成：

$\min_s \sum_{v \in V} s_v^\top (c_v + d_v) + \sum_{(v,u) \in E} s_v^\top R_{vu}\, s_u$

• s_v 是算子 v 的独热选择向量；
• c_v、d_v 是各种切法的通信、计算代价；
• R_vu 是从 v 的第 i 种输出规格到 u 的第 j 种输入规格的重分片代价矩阵。

二次项用标准技巧线性化后交给 ILP 求解器。

几个简化：

• 计算代价设为 0：重算子的每种切法都均分工作，计算量相同；轻量的逐元素算子允许复制计算，代价可以忽略；
• 逐元素、转置等轻量算子合并到它的一个操作数上，缩小 ILP 的规模；
• 解出来以后再做”把 AllReduce 换成 reduce-scatter + all-gather”等优化——达到 ZeRO 式的权重更新分片效果。

5.5 重新发现 Megatron（代码 04 演示 3）

代码的计算图是一条链（QKV → 注意力核心 → 输出投影 → FFN1 → FFN2，不建模残差），链上的这个二次目标可以用逐算子的动态规划精确求解；代码用穷举 26,244 种组合验证了最优值相同。

结果：

• 一个节点的 8 卡上，最优切法是”QKV 按列、注意力按头、输出投影按行、FFN1 按列、FFN2 按行”，重分片代价为 0——正是 Megatron-LM 的切法（第 14 讲）。没有人告诉它 Megatron 的结构，它是从代价里推出来的。最好的逻辑网格是 4 × 2：4 路张量并行 × 2 路数据并行（5.13 ms，纯 8 路张量并行 5.36 ms，纯数据并行 6.92 ms）。
• 2 节点 × 8 卡、一个 micro-batch 4096 token 时，最优是”跨节点 2 路张量并行 + 节点内 8 路数据并行”（通信 2.21 ms；反过来放 2.76 ms）——和”张量并行只放节点内”的经验相反。原因：这里一个 micro-batch 的激活只有几 MB，而一层有 2 亿个参数，跨节点同步权重梯度反而更贵。经验法则依赖于 token 数和权重大小的比例（第 14 讲 b* = n·m/(n + m)）。

6. Alpa（二）：算子间动态规划、运行时、评估 · Alpa: inter-operator parallelism

6.1 目标

把算子序列（按用户定义的顺序）切成 S 级 s₁…s_S，每级分到集群的一个 nᵢ × mᵢ 子网格上，tᵢ 是第 i 级在它的子网格上由算子内 pass 优化后的延迟（前向 + 反向）。B 个 micro-batch 的流水线总延迟：

$T^* = \min_{\text{切分},\ \text{子网格}} \Big\{ \sum_{i=1}^{S} t_i + (B - 1) \cdot \max_{j} t_j \Big\}$

• 第一项：第一个 micro-batch 走过整条流水线的时间；
• 第二项：其余 B − 1 个 micro-batch 由最慢的一级决定。

两个约束：

• 前向算子和它对应的反向算子放在同一个子网格上；
• 子网格必须恰好铺满整个集群，不浪费设备。

子网格形状限制为 (1, 1)、(1, 2)、(1, 4)……(1, M) 和 (2, M)、(3, M)……(N, M)（论文附录证明这些总能铺满 N × M 的集群；被排除的 n > 1、m < M 的形状通常更差）。

6.2 动态规划

先枚举 t_max（单级延迟的上界），对每个 t_max 求”总延迟最小、且每级都不超过 t_max”的切法：

$F(s, k, d; t_{\max}) = \min_{\substack{k \le i \le K,\ n_s m_s \le d \\ t_{\text{intra}}((o_k..o_i), (n_s, m_s), s) \le t_{\max}}} \Big\{ t_{\text{intra}}(\cdot) + F(s - 1, i + 1, d - n_s m_s; t_{\max}) \Big\}$

$T^*(t_{\max}) = \min_s F(s, 0, NM; t_{\max}) + (B - 1) \cdot t_{\max}$

• t_intra 由算子内 pass 给出：对每个候选的逻辑网格形状求一次最优切法，编译、实测，选最快的；
• 还要检查显存：1F1B 下第 s 级（后面还有 s 级）最多同时存 s 份激活，要求 mem_stage + s · mem_act ≤ mem_device。

代码实现了这个动态规划（含显存约束），在 4 层、2 节点 × 2 卡的小问题上和穷举所有”连续切分 + 子网格分配”的结果完全相同。

6.3 结果（代码 04 演示 4–5）

4 节点 × 8 卡、32 个 micro-batch：

• GPT 式：Alpa 把流水线的级放在节点之间（每级一个节点、4 层），级内 4 路张量并行 × 2 路数据并行——就是 Megatron-LM 2021 的手工配置，又是它自己找到的。
• 宽 ResNet 式：前面的层激活大、权重小，用数据并行；后面的层权重大、激活小，切输入通道。这和 slides 的”宽 ResNet 在 16 卡上的切分”案例一致：异构的模型，每一段用不同的并行方式，这种方案很难手工设计。

弱扩展（宽 ResNet 式，每节点的 batch 固定）：从 1 个节点到 4 个节点，Alpa 的每卡吞吐保持在 91%，只用算子内掉到 47%，只用算子间 87%。论文报告宽 ResNet 在 4 个节点上达到 80% 的线性扩展效率（这里的估算更乐观，只看趋势）。

6.4 运行时编排

切好以后还要真的跑起来（slides：“Runtime orchestration”）：

• 跨网格的重分片：相邻两级的网格形状、张量的分片规格可能都不同，这是一个多对多的组播问题。Alpa 先生成点对点的 send/recv，再在目标规格有复制时只跨网格传一份，到达后在目标网格内部用高带宽的 all-gather 分发（“local all-gather”）。
• MPMD 运行时：每一级的算子不同、网格不同，Alpa 为每个网格生成不同的静态指令序列（分配 / 释放显存、跨级通信、同步、计算），执行前一次性下发，避免运行时驱动进程和 worker 之间的协调开销。

6.5 论文的评估

论文在 8 台 AWS p3.16xlarge（64 块 V100）上评估：

• GPT 上 Alpa 和专门的 Megatron-LM 性能相当；
• GShard MoE 上，相比手工调优的 DeepSpeed，2 节点快 3.5 倍、4 节点快 9.7 倍；
• 没有手工方案的宽 ResNet 上，4 个节点达到 80% 的线性扩展效率。

编译和搜索时间随模型大小和卡数线性增长，最大的配置（39B 的 GPT、64 卡）最多几个小时，主要花在枚举并实测”级 × 网格”组合上。

7. 讨论：比较 FlexFlow 和 Alpa · FlexFlow vs Alpa

slides 最后的讨论题。可以从四个角度比较：

一个值得讨论的观点：

• FlexFlow 的空间更大、更”平”（理论上包含 Alpa 能找到的大多数策略），但大空间要靠随机搜索；
• Alpa 用”集群是分层的”这个硬件结构主动缩小空间，换来每一层都能精确求解。

这和第 9 讲 AutoTVM（模板 + 搜索）vs Ansor（自动生成空间）、第 10 讲 TASO（回溯搜索）vs TENSAT（等式饱和）的讨论是同一类问题：空间的大小和搜索的精确性之间的取舍。

Zhihao Jia 组的后续工作 Unity（OSDI 2022）把并行策略和 TASO 式的代数变换放进同一个搜索空间，又往”更大的空间”那边走了一步。

8. 论文精读 · Paper reading

8.1 Beyond Data and Model Parallelism for Deep Neural Networks（FlexFlow，Jia, Zaharia, Aiken, MLSys 2019）

1. 问题：主流框架只支持数据并行（或简单的模型并行），对参数多的算子（全连接、嵌入）效率差。专家设计的混合策略更好，但仍不优，且要为每个模型和集群重做。

2. 为什么困难：

• 策略空间是每个算子配置的笛卡尔积，指数级；
• 在真实硬件上评估一个策略要部署、跑若干次迭代，太慢；
• 计算、通信、负载均衡之间互相牵制，很难写出闭式的代价公式。

3. 以前的工作为什么不行：

• 数据并行、模型并行只是空间里的两个点；
• 专家策略（Krizhevsky 的 one weird trick、GNMT 的层级放置）依赖直觉；
• 基于强化学习的设备放置（REINFORCE）要在真实硬件上评估，要花 12–27 小时、上百个节点，而且只决定放置不切分算子；
• OptCNN 只处理线性的算子图，假设算子不能并行执行，用动态规划。

4. 新意：

• SOAP 搜索空间：把样本、算子、属性、参数四个维度统一起来，每个算子独立选择；
• 执行模拟器：基于”算子性能可预测”实测少量算子，用任务图 + FIFO 模拟，并发明增量模拟让 MCMC 的每一步更便宜；
• MCMC 搜索：把模拟时间当作代价，用 Metropolis–Hastings 在空间里游走；
• FlexFlow 运行时：支持任意维度、算子粒度的并行执行（当时的框架都不支持）。

5. 结果：

• 6 个模型（AlexNet、Inception-v3、ResNet-101、RNNTC、RNNLM、NMT），两个 GPU 集群（P100、K80），比最好的基线快最多 3.8 倍（含搜索时间）；
• 模拟器比真实执行快三个数量级，误差在 30% 以内且保持排序；
• 增量模拟比完整模拟快 2.2–6.9 倍；
• 和 REINFORCE 比，策略快 3.4–3.8 倍，而搜索只要 14–40 秒、一台机器；
• 在 4 块卡的 LeNet、RNNLM 上找到了全局最优。

本讲的验证：在缩小的 SOAP 空间上实现了任务图、完整模拟和 MCMC，AlexNet、DLRM 上找到的策略在结构上符合专家直觉又更进一步；穷举验证了 MCMC 在小问题上能（多数情况下接近）找到最优；但也发现，当好策略由大量小消息组成时，“用满带宽”的假设会让模拟器在好策略之间排错序。

局限：

• 模拟器依赖 A1–A4，动态形状、数据相关的计算不适用；网络拥塞、小消息效率需要单独建模；
• MCMC 没有最优性保证，空间极大时收敛慢；
• 没有流水线（micro-batch）的概念；
• 策略的表达受限于”输出张量怎样切”，无法表达需要改写计算图的优化（Unity 后来解决）。

8.2 Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning（Zheng et al., OSDI 2022）

1. 问题：大模型训练要组合数据、算子（张量）、流水线三种并行，现有系统要么要求用户手写并行方案（Megatron-LM、DeepSpeed），要么只在很小的空间里自动搜索，推广不到异构的模型（MoE、宽 ResNet）和集群。

2. 为什么困难：

• 三种并行的组合空间巨大，而且互相影响：流水线的切分决定了每级的子图和子网格，子网格的形状又决定了级内能用哪些切法；
• 集群是异构的（节点内外带宽差一个数量级）；
• 每个算子的切法之间还有重分片的代价。

3. 以前的工作为什么不行：

• 手工系统（Megatron-LM、DeepSpeed 3D 并行）只对特定模型调优；
• FlexFlow、Tofu 等自动方法只覆盖算子内（或扁平的空间），不显式建模流水线；
• PipeDream、DAPPLE 只做流水线切分，级内只用数据并行；
• GSPMD 需要用户标注分片。

4. 新意：

• 重新分类：按”是否切开算子”把并行分成算子间 / 算子内两层，并把它们映射到集群的两层带宽结构上；
• 算子内 pass：分片规格 + 重分片代价 + ILP（节点代价 + 边上的重分片代价）；
• 算子间 pass：动态规划同时决定流水线切分和子网格分配，每个候选由算子内 pass 优化后编译实测，并检查 1F1B 的显存约束；
• 运行时编排：跨网格的重分片（local all-gather）、MPMD 静态指令；
• 作为 JAX + XLA 的编译器实现，用户只需加一个装饰器。

5. 结果：

• GPT 上和 Megatron-LM 相当（“在它擅长的模型上打平专家”）；
• MoE 上比 DeepSpeed 快 3.5 倍（2 节点）、9.7 倍（4 节点）；
• 宽 ResNet 4 节点 80% 的线性扩展效率；
• 编译最多几个小时。

本讲的验证：缩小版在 GPT 式模型上重新发现了 Megatron 的层内切法和”流水线跨节点、张量并行在节点内”的配置；在宽 ResNet 式模型上前段数据并行、后段切通道；两层的求解都用穷举验证过最优；也发现在小 micro-batch 时，算子内 pass 会把张量并行放到跨节点的轴上——经验法则不是普遍成立的。

局限（论文 §7）：

• 不建模级间通信的代价（假设它很小）；
• micro-batch 数 B 是超参数，要枚举；
• 只考虑静态的线性流水线，不考虑不同分支在不同设备上并行；
• 不优化计算和通信的重叠；
• 只处理形状在编译时已知的静态图；
• 两层分解本身是一种近似：算子间 pass 只能看到算子内 pass 报告的延迟，不能联合优化。

8.3 两篇论文之间的联系

• 同一个范式：都是”搜索空间 + 代价模型 + 搜索算法”，都用”实测少量东西 + 模型组合”代替在硬件上评估每个策略（FlexFlow 实测算子，Alpa 实测编译后的级）。
• 对空间的不同处理：
  • FlexFlow 用一个统一、扁平的空间，靠随机搜索；
  • Alpa 利用硬件的层次结构把空间分解，每层用精确的组合优化（ILP、DP）。
  • Alpa 论文把 FlexFlow 归为”在有限空间里搜索”（不含流水线）；反过来，Alpa 也失去了 FlexFlow 的”不同分支同时在不同设备上跑”。
• 代价模型的假设决定结论：本讲的 02 号脚本（FlexFlow 模拟器在小消息上排错序）和第 14 讲的 PipeDream 切分（论文式子 vs 环形 AllReduce）讲的是同一件事——自动化系统的结论只在它的代价模型准确时才可信。
• 一个可以讨论的结合：用 Alpa 的两层分解缩小空间，再在每层用 FlexFlow 式的模拟器（而不是只算字节数）评估重叠和争用；或者反过来，把 Alpa 的流水线维度加进 SOAP。

9. 小结 · Summary

1. 为什么要自动并行：组合数据、张量、流水线并行的空间太大，手工设计难、不优、不可移植。自动并行 = 搜索空间 + 代价模型 + 搜索算法。
2. SOAP：每个算子沿样本、属性、参数维切，加上算子之间的设备分配。同一个计算的不同切法结果相同、代价不同（同步参数 vs 搬输入 vs 光晕）。空间大到 10¹¹–10²² 以上。
3. 执行模拟器：任务图（计算任务 + 通信任务）+ FIFO 的完整模拟；增量模拟只重算变化的部分。它依赖”算子时间可预测、通信能用满带宽”等假设；保持好策略之间的排序比平均误差更重要，而搜索恰恰会走向模型最乐观的地方。
4. MCMC：随机改一个算子的配置，按 exp(β·Δcost) 接受。AlexNet 上比数据并行快 16.9 倍、比专家快 2 倍，DLRM 上嵌入表不再复制；小问题上能找到或接近全局最优；β 太小乱走、太大贪心。
5. Alpa 算子内：分片规格（R / S0 / S1 / S01）、重分片（本地切 / all-gather / all-to-all）、每个算子的切法、ILP 目标”节点代价 + 边代价”。它能从代价里重新推出 Megatron 的切法。
6. Alpa 算子间：动态规划同时切流水线和分子网格，T* = Σt + (B − 1)·max t，含显存约束。GPT 上得到”流水线跨节点、张量并行在节点内”，宽 ResNet 上得到”前段数据并行、后段切通道”。
7. FlexFlow vs Alpa：扁平的大空间 + 随机搜索 vs 分层的空间 + 精确求解；共同的教训是代价模型决定结论。

10. 常见误区 · Common misconceptions

• “自动并行能找到全局最优。” 一般不能：问题是 NP 难的。FlexFlow 的 MCMC 没有保证；Alpa 的每一层在它的代价模型下是精确的，但两层分解本身是近似，代价模型也有误差。
• “模拟器误差在 30% 以内就够用了。” 不一定。代码里 DLRM 的整体误差和排序都”还行”，但最好的那些策略之间排序是乱的，模拟器挑的”最好”策略实际排第 48。
• “SOAP 的属性维和样本维一样。” 都复制参数，但属性维切开后，卷积之类的算子要多读边界上的”光晕”，而且属性维可以在 batch 很小时提供并行度（例如大图像、长序列、图计算）。
• “Alpa 的算子内并行就是张量并行。” 算子内并行包含数据并行（把轴分给 batch 维）、张量并行（分给特征维）以及两者的组合，ZeRO 式的权重更新分片也可以由它推出。
• “张量并行一定要放节点内。” 通常如此，但取决于激活和权重的相对大小。代码里 micro-batch 只有 4096 个 token 时，最优解把 2 路张量并行放在了跨节点的轴上。
• “有了自动并行，就不需要理解手工策略了。” 恰恰相反：读懂自动系统给出的方案（为什么卷积数据并行、为什么流水线跨节点），要靠第 12–14 讲的手工分析；而判断它的代价模型在你的场景里是否成立，更需要这些知识。

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

练习 1　一个 Linear 层 Y = X·W（X：b × n，W：n × m），在 SOAP 里沿样本维切 4 份、沿参数维切 2 份，放在 8 块卡上。每块卡要读多少输入、存多少权重？迭代末尾要同步什么？

答案：每块卡负责 b/4 个样本、m/2 个输出通道。读输入 (b/4) × n，存权重 n × (m/2)。同一个参数分片（同一个通道下标）有 4 个副本（样本维的 4 份），迭代末尾在每组 4 块卡之间 AllReduce n·m/2 个梯度元素。

练习 2　FlexFlow 的接受准则是 min(1, exp(β·(cost(S) − cost(S*))))。设当前策略 100 ms，提议的策略 110 ms。β = 0.1 / ms 时接受概率是多少？β = 1 / ms 呢？为什么 β 的合适取值和代价的尺度有关？

答案：

• β = 0.1：exp(0.1 × (100 − 110)) = exp(−1) ≈ 0.37；
• β = 1：exp(−10) ≈ 4.5 × 10⁻⁵，几乎不会接受。
• 如果代价的单位从 ms 换成 s，同样的 β 意味着完全不同的接受率。代码因此把代价差除以初始代价，让 β 无量纲。

练习 3　FlexFlow 的增量模拟为什么和完整模拟的结果完全相同？什么时候它的加速最小？

答案：增量模拟从旧的时间线出发，把被改动的任务入队，像 Bellman–Ford 一样只要某个任务的就绪或开始时间变了就继续向后传播，直到不再变化。最终的固定点和从头模拟的结果相同（FIFO 顺序也被维护）。如果改动的算子在关键路径的很前面、影响了几乎所有后续任务，就要重算大部分时间线，加速最小；卡越多、图越大，一次改动影响的比例越小，加速越大（论文 2.2 → 6.9 倍）。

练习 4　写出把规格 S0S1 的张量变成 RR 需要的通信（2 × 4 网格）。哪一步更贵？

答案：两个轴都要 all-gather：沿轴 1 all-gather（4 块卡），沿轴 0 all-gather（2 块卡）。先做哪个都可以，但每一步的数据量不同：先沿轴 1 做（快的轴），每卡数据从 1/8 变成 1/2，再沿轴 0（慢的轴）传 1/2 的数据；反过来先沿轴 0 只传 1/8，再沿轴 1 传 1/4……总量相同，但慢的轴上传的量不同——先在慢轴上做、趁数据还小，通常更便宜。（代码里的规则按轴 0、轴 1 的固定顺序做，是一个简化。）

练习 5　Alpa 算子内的 ILP 目标为什么有二次项 s_vᵀ R_vu s_u？在链式图上为什么可以用动态规划精确求解？

答案：重分片代价取决于两个相邻算子各自选了什么切法，是两个选择变量的乘积。在链上，第 k 个算子的最优累计代价只依赖第 k − 1 个算子的选择（马尔可夫性），所以”前 k 个算子、第 k 个选 a 的最小代价” = min over 第 k − 1 个的选择 b {前 k − 1 个选 b 的最小代价 + R_ba} + c_a，逐个推下去就是精确解（Viterbi）。一般的图有分支和汇合，不再有这个结构，所以用 ILP。

练习 6　Alpa 的算子间目标 Σtᵢ + (B − 1)·max tⱼ 为什么要先枚举 t_max？直接对整个式子做动态规划会有什么问题？

答案：max 项依赖所有级的选择，不能分解成”前面的最优 + 后面的最优”——一个总延迟更小的前缀可能有更大的单级延迟，从而在 max 项上吃亏。枚举 t_max 后，max 项变成常数 (B − 1)·t_max，约束”每级 ≤ t_max”可以逐级检查，剩下的 Σtᵢ 是可加的，就能动态规划。最后在所有 t_max 里取最小（代码还用”(B − 1)·t_max 已经超过当前最好解”提前停止）。

练习 7　为什么 DLRM 上数据并行特别差？SOAP 里怎样的配置能避免这个问题？

答案：嵌入表的参数极多（代码里 8 张表共 5.12 亿个参数），但每个样本只查每张表的一行，计算量很小。数据并行要在每步对整张表的梯度做 AllReduce（代码里 2 GB），通信完全压倒计算。SOAP 里给嵌入算子选”不切 + 放在一块卡上”（算子维）或者”沿参数维切”（嵌入维度），不复制参数，只把查出来的向量传给后面的 MLP。

练习 8　（开放）FlexFlow 的模拟器能模拟计算和通信的重叠，Alpa 的算子内代价只算通信字节数、不计计算。各有什么好处和代价？如果要在 Alpa 里加入重叠，你会怎么做？

参考思路：

• 模拟器更精确（重叠、争用、关键路径），但每次评估慢，而且不能放进 ILP 这样的精确求解器（代价不是简单的加和）。
• Alpa 的加性代价让 ILP 可解、可以精确求最优，但会高估不重叠的通信、忽略关键路径。
• 折中：在 ILP 里只粗略地建模，解出若干个候选，再用模拟器或实测精排（第 9 讲 Ansor 的”代价模型筛选 + 实测”思路）；或者在算子间 pass 的实测里自然地包含重叠（Alpa 已经对每个级实测）。

12. 延伸阅读 · Further reading

• C. Unger et al., Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization（OSDI 2022）——FlexFlow 的后继，把 TASO 式的图替换和并行一起搜索。
• Z. Jia et al., Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks（OptCNN，ICML 2018）——层级并行的动态规划。
• Y. Xu et al., GSPMD: General and Scalable Parallelization for ML Computation Graphs（2021）——用分片标注 + 传播实现 SPMD 并行；Alpa 定好切法以后，用 XLA 和 GSPMD 编译每一级的并行可执行文件。
• M. Wang, C. Huang, J. Li, Supporting Very Large Models using Automatic Dataflow Graph Partitioning（Tofu，EuroSys 2019）。
• S. Fan et al., DAPPLE: A Pipelined Data Parallel Approach for Training Large Models（PPoPP 2021）。
• A. Mirhoseini et al., Device Placement Optimization with Reinforcement Learning（ICML 2017）——FlexFlow 比较的 REINFORCE 基线。
• PyTorch 的 DTensor / torch.distributed.tensor.parallel：用 Shard / Replicate / Partial 描述分片，和 Alpa 的规格一一对应。

13. 下一讲预告 · Next lecture

第 12–15 讲解决的是训练的并行。第 16 讲转向服务：一个训练好的 LLM 怎样又快又省地回答成千上万个请求？

• 自回归解码分预填充和解码两个阶段，解码时每一步只处理一个 token，GPU 大部分时间在等显存；
• 连续批处理（Orca）在每一次迭代而不是每一个请求的粒度上调度；
• PagedAttention（vLLM）像操作系统的分页一样管理 KV 缓存，把显存浪费从 60–80% 降到不到 4%；
• RadixAttention（SGLang）用基数树在请求之间共享前缀的 KV 缓存，并用重叠调度藏住 CPU 开销。

我们会写一个离散事件的服务模拟器、一个分页的 KV 分配器和一棵带 LRU 淘汰的基数树，量一量它们各自省了多少。