预训练阶段。让模型在超大规模的通用数据集上进行无监督或者自监督训练，以掌握基础的语言规律、世界知识或图像特征的过程。
关于集体通信
集体通信 collective communication 是分布式训练中一系列 GPU 间的数据操作。
All-Reduce
1 | Before: |
每张 GPU 的对应元素进行 reduction，然后每张 GPU 都得到完整结果。适用于梯度同步等场合。
All-Gather
1 | Before: |
每张 GPU 原本只保存部分数据，经过 All-Gather 后，所有 GPU 都得到完整数据。ZeRO-3 临时获取完整参数的时候有用到。
Reduce-Scatter
1 | Before: |
先做 reduction，然后把结果切分给不同 GPU。ZeRO 梯度同步的时候用到。
All-to-All
1 | Before: |
每张 GPU 把自己的数据切分成若干块，分别发送给其他 GPU，同时接收其他 GPU 发来的数据块。
DP，TP 和 PP
总之就是模型越来越大，预训练的时候单个 GPU 显存不够，要找各种方法把超级大的 Transformer 装到一堆 GPU 上。
Data Parallelism
DP 的话就是每张 GPU 放一份完整模型，把训练数据切成多个 mini batch，每个 GPU 负责一个，然后同步 gradient（叫做 all-reduce gradients）。
Model Parallelism: TP
但是模型已经大到一张 GPU 已经装不下了，所以要想办法把模型本身也给切了，分到多个 GPU 上。切模型就是 Model Parallelism（MP）。
有两种切法，一种是在 layer 里面内切，叫做
矩阵乘法 $Y=XW$。把参数纵切的话就是 column parallel linear $W=[W_0, W_1]$，GPU 0 负责 $Y_0=XW_0$，GPU 1 负责 $Y_1=XW_1$，最后 all-reduce 得到 $Y=[Y_0, Y_1]$。
也可以横切 $W=
\begin{bmatrix}
W_0\\
W_1
\end{bmatrix}
$，把输入切成 $X=[X_0, X_1]$，最后加起来 $Y=X_0W_0+X_1W_1$，这是 row parallel linear。
如果在一些机器学习框架中看到参数 TP=8 的话，意思就是把 layer 切成 8 份。
不过 TP 有个问题就是需要频繁的 GPU 间通信，毕竟每一层都被切成几份放在不同 GPU 上，做 all-reduce、all-gather、reduce-scatter 这些拼结果的操作都依赖通信。所以一个实用的原则是做 TP 的训练尽量放在一个 node 里面（一般同一个 node 里的 GPU 有高速通信比如 NVLink）。
Model Parallelism: PP
另一种切法是按 layer 来切，每张 GPU 负责若干层 layers，这叫做
就像这样：
1 | time → |
GPU idle 的区域被称为
1F1B Megatron
Deepak Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM,” SC 2021.
这篇论文干了什么呢，就是设计了一个算法，把 DP、TP 和 PP 全用上，也就是所谓的 3D Parallelism。并且还研究了各种 pipeline 的调度，提出了 interleaved（一张 GPU 负责不连续的 layer chunk）1F1B（forward batch 和 backward batch 交替处理）调度算法，显著缩减了 pipeline bubble，提高 throughput。他们试过最后在 3072 张 A100 上训了一个 1T 的模型（壕。。。）。
ZeRO
想起三年前被 ZeRO 和 OOM 支配的恐惧。。。
Samyam Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” SC 2020.
ZeRO 主要是针对 DP 的。DP 的话，每张 GPU 都要拥有一份模型状态的完整拷贝，实在是太浪费了。
ZeRO-1 先把 optimizer states（比如说 Adam 的 moment）给切了，每张 GPU 只存 $1/N$。根据一个神奇公式推出来 $N$ 很大的时候 ZeRO-1 可以省 4x 显存。
ZeRO-2 进一步把 gradients 也给切了。它可以省 8x 显存。ZeRO-1 和 ZeRO-2 可以做到和普通 DP 大体相同的通信 volume。
ZeRO-3 再把参数也给切了。这下模型状态的所有东西都被切成了 $N$ 份，理论上 GPU 数目增加 $N$ 倍，显存就可以节省 $N$ 倍。
那么有人会问了，那都是切模型，为什么 ZeRO-3 属于 DP，而 TP 属于 MP 呢？因为 TP 的话，一个矩阵乘法是真的被分给多个 GPU 共同计算的，而 ZeRO-3 只是在切模型状态的存储，真的在算 layer 的时候，每张 GPU 会通过 all-gather 暂时得到完整的模型状态，各自做各自的计算。简单来说，ZeRO-3 只切 storage，而 TP 把 computation 也给切了。
ZeRO-3 还设计了通信的调度，据说只会增加大约一半的通信 volume。
FSDP
GPipe
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 都得到了利用。
1 | hidden H |
总之，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 可能需要进行各种 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 了。
SP
Sam Ade Jacobs et al., DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models, 2023.
这篇处理的是 SP，即 Sequence parallelism。
DP 将一个 batch 中的不同 samples 分给不同 GPU 并行处理，但每张 GPU 仍然需要处理完整的单条 sample。对于长序列训练，即使 micro-batch size 已经是 1，单条 sample 仍可能因为 sequence length 太长而无法高效处理。
SP 针对的就是长序列样本的并发。很自然地，只需要把长序列切开平分给每张 GPU 不就行了嘛！假如说有四张 GPU，一条 128K token 的长序列，每张分别负责 32K token。但一个问题是，global attention 这个操作要求每个 token 都 attend 整个序列中的其他 token。
DeepSpeed-Ulysses 通过两次 All-to-All，既保持了 sequence sharding，又让分布在不同 GPU 上的序列 token 做完整的 global attention。
第一次 All-to-All 很巧妙，它把数据按 sequence 切分转换成按注意力头切分。比如说，原来 GPU0 负责 0-32K token 的 H0-3 四个注意力头的计算，转换成负责 0-128K token 的 H0 单个注意力头的计算；其他三张 GPU 也做类似的转换。注意到转换前后每张 GPU 持有的 QKV 数据量没有变化！既然每张 GPU 现在都有了整条序列，而注意力头之间的计算是独立的，直接算 attention 就可以了。第二次 All-to-All 就把数据交换回来，恢复原来的 sequence partition，以便 output projection、MLP、LayerNorm 这些操作能正常地按照 SP 进行。
之前的 SP 方法比如 Megatron-SP，主要使用 All-Gather/Reduce-Scatter 在 SP 和 TP 之间转换，通信量随序列长度线性增长。DeepSpeed-Ulysses 则利用 sequence-to-head 的 trick，将通信量降低到 $O(Nh/P)$，其中 $h$ 是 hidden size，$P$ 是 SP degree。所以如果序列长度和 GPU 数量同比例增加，通信量就可以保持近似不变，从而实现更好的 scalability。