这一节主要讲的是训练模型的时候的一些并行化的手段

并行化的原因：单个GPU的算力/内存都有限。需要更多的节点来扩展算力，并把模型放入到内存中。
我们的核心目的（也是并行化的核心目的），就是线性的scale。随着GPU数量上升，max model params和flops都可以线性上升。

GPT-NeoX-20B: An Open-Source Autoregressive Language Model 中非常好的一张图，展示了GPU并行训练时候的关键节点和通信链路。
* GPU/GPU通信可以走nvlink，带宽非常高
* GPU/CPU需要走PCIE，相对比较慢
* 节点之间的通信可以走InfiniBand

GPU的互联有上限，最多256个进行互联。

一些基础的通信算子。主要看后面的两个：
* All Gather：相当于把每个分片的slice广播到其他的分片
* Reduce Scatter：分片做reduce

一个小点，就是All reduce等于先做Reduce Scatter，再做All Gather。这样拆成两个算子可以让我们在中间插入一些自己的计算逻辑。

然后开始讲并行化的各种方法：
* Data Parallelism
* Naive data parallel
* ZeRO
* Model Parallelism
* Pipeline parallel
* Tensor parallel
* Activation Parallelism
* Sequence parallel

Naive data parallelism

最简单的并行化方法，一批数据分开放到若干个GPU中，分别计算，再累计梯度。
* 每次需要做一下参数的All Reduce——把参数分发出去，然后收集梯度。所以通信开销是2x params count

这个方法的问题就是我们没有scale memory，模型过大的时候无法放到GPU内存中。

分析内存占用可以发现，大部分的内存占用在optimizer state中：
* 因为optimizer state是持续累积的，而不是像梯度/激活值一样每次计算，所以需要用精度高一些的类型来存储数据。

ZeRO

有了比较经典的方法ZeRO来优化内存占用，从最大头的optimizer state开始shard，然后是gradient，最后则是parameter

最简单的ZeRO1:
* 每一个GPU worker保留optimizer state的一个shard
* 反向传播后，每个worker有自己这一批batch的梯度，通过reduce scatter，每个worker去收集自己对应optimizer shard的梯度。
* 更新optimizer。这里的核心假设是optimizer状态相互不依赖。（至少AdamW是没有依赖的）
* optimizer更新参数，再把参数做All gather广播出去。

通信开销这里可能还得单独看看，这个再去读一下ZeRO的论文吧，因为通信开销和网络拓扑也有关系。

直观讲和Naive DDP相比传输的数据少了。（不过这个应该也看all reduce的实现）。在没有开销的情况下实现了optimizer state的shard。让部分参数可以线性扩展了。

ZeRO2:
* shard gradient。这块的难点主要在每个机器上是不同的activation，然后gradient的计算是有依赖的。
* 同步的进行反向传播，当计算完gradient之后，立刻把gradient发送给对应的worker。
* 相当于在state1的基础上，快速做一下gradient的reduce scatter。
* 因为state1上每一个worker只收集自己shard的gradient，所以收集时机可以提前，从而降低对保留gradient这段内存的要求。

ZeRO3，也就是知名的FSDP(full shard data parallel)：
* param做shard的难点是，data parallel需要在每个机器上做计算，但是如果只有一个shard有param，就需要在计算的时候把param广播出去。
* 反向传播的时候也需要param参与前面梯度的计算，所以反向传播的时候也需要把param再gather一把
* 通过通信计算overlap来降低对延迟的影响

pytorch这里也有一个相关论文：PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

然后看一下图上的例子：
* AG0先广播param，然后计算activation，得到FWD0。
* 在计算的同时，开始广播下一层AG1。然后立刻计算FWD1。
* 没有依赖的param就可以释放了（这里其实感觉可以做的比较多？因为后续反向传播也可以用这里的weight，所以越靠后这个param就越不应该释放，用memory去trade 带宽）
* 反向传播的时候也是一样的。通过All gather拿到param，然后计算梯度。计算梯度后可以通过reduce scatter把梯度发送给对应的shard。此时梯度就可以立刻释放了。

ZeRO的好处是概念简单，和模型没有耦合。比如在每一层上实现一个FSDP wrapper，就可以实现上面的事情，直观想：
* 可能有一个调度器，去做每一层的param/gradient的prefetch。以及对应这一层的计算。
* 不知道是否可以把这里的通信逻辑放到一个算子里，那么这里就变成一个计算/通信的DAG，按顺序执行即可。

这里的data parallel的方法可以让我们把模型做分片，有两个问题：

* 需要比较大的batch size来提高计算强度。否则的话瓶颈主要在通信开销上。(paper: An Empirical Model of Large-Batch Training)
* 但是batch size无法非常大，需要作为一种有限的资源来分配。

直观理解一下这里的batch size的上限问题：
* 这里并不是说batch变大不好，而是说scale的不好
* 随着batch size的变大，对训练效率的提升会逐渐收敛。单个样本的贡献会逐渐变小。
* 此时再去加硬件来提高batch size，带来的收益不一定比较高。此时可能把硬件放到其他的地方（比如提高单个step的速度）会更好。

第二个问题就是ZeRO没有shard activation，所以activation（batch size仍然是有限的）

Pipeline Parallel

因为深度神经网络是一层一层堆出来的，所以一个直观的切分方式就是按层切

* 直观问题就是利用率太低了，一个batch的数据同一时间只在一个GPU上做运算。

CPU这边已经有很多相关的经验了，做这个pipeline的时候需要切stage，一个worker只负责一个stage，这样理论上就可以有stage个worker并发了。

这里有一个简单的并行化的方法，就是batch size内去切。可以减少这个bubble的大小。
* 但是bubble的比例和batch size是有关的。如果batch size很小，就退化到上面的case。
* bubble的大小和三角形的面积相关。算出来应该是stage_size * stage_size / batch_size * stage_size

Pipeline parallel的好处：
* 可以把layer分发到不同GPU上，更好的做memory的scale
* 通信开销比较低，只需要传播activation。而且不需要broadcast，只需要点对点。
* 所以一般适用于比较慢的网络。比如nvlink之间的node，用pipeline。
缺点：
* 依赖batch size去隐藏bubble

还有一些更激进的pipeline parallel的手段，不过感觉整体思路应该离不开CPU那些，毕竟是相似的设计，相似的思路应该都可以从CPU哪里借鉴到。

Tensor Parallel

TensorParallel的思路就是切分矩阵乘法。拆成若干个子的矩阵进行相乘，可以利用多个GPU的算力。

直观的想，需要引入更多的通信开销，因为不是单纯去广播activation/param，需要把矩阵都广播出去，然后做reduce

思路上也比较简单，因为可以做到模型层，然后把不同的部分放到不同的GPU中（类似专家并行）

比如上面这个case，不需要做完XA之后立刻reduce，因为后面的GeLU等操作都是没有相互依赖的，就可以尽量延迟reduce的操作

节点内的parallelism比较有效，否则就会成为memory bound。从图中可以看到，8 -> 16的scale，GPU算力下降了一半，说明扩展到16node的算力和8node是一样的，就是无效的扩展。

和Pipeline parallel的对比：
* 复杂度比较低（不需要做stage的拆分，可以做到模型层）
* 也不需要很大的batch size，比如param大直接拆param就行，一样可以利用算力。
缺点：
* pipeline只需要传递activation
* tensor parallel需要传递activation，param等，并且每次计算需要一次广播，然后计算完可能会再reduce（不过这个8倍不太清楚是怎么算出来的）

所以tensor parallel适合节点内并行，pipeline parallel适合跨节点的

Activation memory

最后一块是考虑activation memory的切分。

上面的ZeRO是把param，grad，optimizer state做shard。pipeline则是根据层切。无法扩展activation memory。（比如batch size变大的时候无法处理）

tensor parallel好像可以处理，但是如果不把activation放到GPU的话，会持续的在DRAM中换来换去

activation memory的占用。后面有一个根据sequence length的二次占比的就是attention层。flash attention可以去掉这一项

用tensor parallel的话（不知道考虑的是峰值还是什么，峰值的话如果全在GPU还是需要all gather的），可以都除t。

34中，有24是矩阵乘法相关，可以被tensor parallel扩展。另外的操作是pointwise的操作，比如layer norm。无法通过tensor parallel扩展。

但是，因为他们是pointwise的，所以可以不用管sequence这个属性，直接分发给所有的GPU做pointwise的操作。比如这里layer norm是pointwise的，在每个GPU上执行完之后，去all gather。然后放到self attention中（因为这里需要根据sequence做聚合了）。