自动并行 — Lalande.

Alpa

你在这篇文章里加了什么。。我读得好累。不愧是 OSDI 啊。

Lianmin Zheng et al., “Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning,” OSDI 2022.

Alpa 把 parallelism 建模成一个巨大的组合优化问题，并且设计了一个编译器进行求解。

Intra-operator Parallelism

即一个 operator 本身被分给不同 GPU 进行计算。比如说矩阵乘法 $C=AB$ ，可以把 $A$ 、 $B$ 沿某个 dimension 切开，分给四张 GPU 一起算。DP（切 $B$ ），TP（切 $A$ 、 $B$ ），ZeRO 等 tensor sharding 的方式都属于 Intra-operator Parallelism。

Inter-operator Parallelism

即不同的 operators 如何分给不同的 GPU groups（或者称为 stages）。所以 PP 属于 Inter-operator Parallelism。

Device Mesh

为方便起见，把 GPU 集群组织成二维的逻辑空间 device mesh。因此，mesh axis 自动定义了 GPU 的通信分组方式。为什么是二维呢？一个很自然的原因是 GPU 集群通常具有两层通信拓扑，例如节点 × 同节点内 GPU：同一节点内的 GPU 通信较快，而跨节点通信较慢。所以，device mesh 不仅描述了 GPU 如何分组，也可以反映不同 axis 上的通信代价。

考虑一个 $2\times 4$ 的 device mesh 与 $B\times H$ （batch 和 hidden 维度）的 tensor $X$ 。假如指定 $B$ 沿 axis 0， $H$ 沿 axis 1 做 sharding，这就代表着 DP = 2，TP = 4，且 8 张 GPU 都得到了利用。

code

                 hidden H
              ─────────────→
             H0   H1   H2   H3
           ┌────┬────┬────┬────┐
batch B0   │ G0 │ G1 │ G2 │ G3 │
           ├────┼────┼────┼────┤
batch B1   │ G4 │ G5 │ G6 │ G7 │
           └────┴────┴────┴────┘

总之，mesh 的维度代表设计者想为可并行的计算提供多少个独立的 device 维度，和每个 device 维度可以提供多大的 parallel degree。

Sharding Specs

Alpa 的统一 sharding 框架。我们不再关注到底是 DP、TP 还是 ZeRO，只用知道它们本质上都是一个将 tensor axes 映射到 mesh axes 上的过程。这被称为 sharding specs，或 sharding layout。他们使用的 notation 是这样的：

$R$ ：replicated
$S^0$ ：沿 mesh axis 0 做 sharding
$S^1$ ：沿 mesh axis 1 做 sharding
$S^{01}$ ：同时沿 mesh axis 0 和 1 做 sharding

考虑一个线性层 $Y=XW$ ，其中 $X[B, H]$ ， $W[H, O]$ ， $Y[B, O]$ ，我们分别来看看 DP、TP 和 ZeRO 是怎么表示的。

DP 的 $X$ 是 $S^1 R$ （输入数据按照 batch 切）， $W$ 则是 $RR$ （每张 GPU 保存完整参数）， $Y$ 是 $S^1 R$ （输出仍然按照 batch 切）。TP 的话，以 column parallelism 为例， $X$ 当然是 $RR$ （输入数据完整复制），参数 $W$ 则是 $RS^1$ （column 的话切的是输出维度）， $Y$ 对应的是 $RS^1$ 。ZeRO 同理，只不过我们需要进一步考虑梯度、optimization states 这些 tensor 的切法。

不同的 operator 最适合的 tensor sharding strategies 可能是不同的。例如上一个 operator 输出的 tensor 是沿 batch 维度切的，但下一个 operator 希望它沿 hidden dimension 来切，那么中间就必须重新分布数据，这个过程叫做 resharding。

resharding 可能需要进行各种 collective communication，代价是比较昂贵的。

问题现在变成这样了：给定一个 stage 和 device mesh，对于每个 operator，存在若干可能的 sharding strategies，它们的（自身与和相邻 operator 间可能的 resharding）计算、通信和空间代价是可估计的。这不是一个全局的组合优化问题的形式吗？这样，就可以用整数线性规划 ILP 来求解了。

Alpa compiler

用户只需要写 model，Alpa 考虑的就多了。

Alpa compiler 会分析 model 的 computation graph，建立 device mesh，生成可能的 sharding strategies，利用 ILP 搜索较好的方案，并自动插入必要的 All-Reduce、All-Gather、All-to-All 等通信操作。这样，intra-op parallelism 就在 compiler 的层面被解决了。

至于 inter-op parallelism，即 computation graph 怎么划分成 pipeline stages，以及每个 stage 分配多少 GPU，这是一个依赖于 intra-op parallelism 的问题。在 ILP 求解出不同 stage 和 device mesh 组合估计的执行代价后，就可以通过动态规划寻找整体较好的 pipeline plan 了。