并行训练

预训练阶段。让模型在超大规模的通用数据集上进行无监督或者自监督训练，以掌握基础的语言规律、世界知识或图像特征的过程。

关于集体通信

集体通信 collective communication 是分布式训练中一系列 GPU 间的数据操作。

All-Reduce

code
Before:
GPU0: [1, 2]
GPU1: [3, 4]

After:
GPU0: [4, 6]
GPU1: [4, 6]

每张 GPU 的对应元素进行 reduction，然后每张 GPU 都得到完整结果。适用于梯度同步等场合。

All-Gather

code
Before:
GPU0: [A]
GPU1: [B]

After:
GPU0: [A, B]
GPU1: [A, B]

每张 GPU 原本只保存部分数据，经过 All-Gather 后，所有 GPU 都得到完整数据。ZeRO-3 临时获取完整参数的时候有用到。

Reduce-Scatter

code
Before:
GPU0: [1, 2]
GPU1: [3, 4]

Reduce:
      [4, 6]

After:
GPU0: [4]
GPU1: [6]

先做 reduction，然后把结果切分给不同 GPU。ZeRO 梯度同步的时候用到。

All-to-All

code
Before:
GPU0: [A0, A1]
GPU1: [B0, B1]

After:
GPU0: [A0, B0]
GPU1: [A1, B1]

每张 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 里面内切，叫做 Tensor Parallelism（TP）：

矩阵乘法 Y=XWY=XW。把参数纵切的话就是 column parallel linear W=[W0,W1]W=[W_0, W_1]，GPU 0 负责 Y0=XW0Y_0=XW_0，GPU 1 负责 Y1=XW1Y_1=XW_1，最后 all-reduce 得到 Y=[Y0,Y1]Y=[Y_0, Y_1]。

也可以横切 W=[W0W1]W=\begin{bmatrix}W_0\\W_1\end{bmatrix}，把输入切成 X=[X0,X1]X=[X_0, X_1]，最后加起来 Y=X0W0+X1W1Y=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，这叫做 Pipeline Parallelism（PP）。为什么是 pipeline 呢，因为训练数据是按顺序经过一层层 layer 的，把数据再切成 mini batch，就很像不断在传送带上传送东西。

就像这样：

code
time →

GPU0: M1 M2 M3 M4 M5 ...
GPU1:    M1 M2 M3 M4 ...
GPU2:       M1 M2 M3 ...
GPU3:          M1 M2 ...

GPU idle 的区域被称为 pipeline bubble。

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 的模型（壕。。。）。

GPipe

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/N1/N。根据一个神奇公式推出来 NN 很大的时候 ZeRO-1 可以省 4x 显存。

ZeRO-2 进一步把 gradients 也给切了。它可以省 8x 显存。ZeRO-1 和 ZeRO-2 可以做到和普通 DP 大体相同的通信 volume。

ZeRO-3 再把参数也给切了。这下模型状态的所有东西都被切成了 NN 份，理论上 GPU 数目增加 NN 倍，显存就可以节省 NN 倍。

那么有人会问了，那都是切模型，为什么 ZeRO-3 属于 DP，而 TP 属于 MP 呢？因为 TP 的话，一个矩阵乘法是真的被分给多个 GPU 共同计算的，而 ZeRO-3 只是在切模型状态的存储，真的在算 layer 的时候，每张 GPU 会通过 all-gather 暂时得到完整的模型状态，各自做各自的计算。简单来说，ZeRO-3 只切 storage，而 TP 把 computation 也给切了。

ZeRO-3 还设计了通信的调度，据说只会增加大约一半的通信 volume。

FSDP