因为写代码鸽了一段时间…接下来会开始逐步整理这段时间积累的一些经验。
在开始介绍并行技术之前，有必要先介绍一下我们的优化目标：
- 能训练更大的模型。期望可训练模型的大小和GPU的数量呈线性关系。
- 能更快的训练模型。期望训练模型的速度和GPU的数量呈线性关系。

模型并行

如果一张卡装不下这个模型。
由于神经网络的层级思想，我们可以把神经网络分为不同的块。按照顺序放置在指定的GPU上。
大概就是这样。

时间上的问题

所以现在进行一轮训练，大概是如下图这样

每一个行代表在GPU上的独立的神经网络块。按照顺序，依次执行Forward操作和Backward的操作。
但是这样很显然，GPU的空转问题很严重。甚至可以说大部分的时间，GPU都是在空转的。假设有块GPU，那么我们很容易可以退出，实际无用的时间比例为：
这样，K越大的时候，无用的时间就越多。GPU的资源都被浪费了，因此我们需要解决这个问题。

空间上的问题

在进行Backward的计算过程中，我们假设batch_size大小为N，模型有L层，每一层储存的激活值占用的空间平均为d，那么对于每块GPU，其额外占用的空间为。从这个复杂度可以看出，随着三者的增加，K的增加可能就微不足道了。因此我们需要想办法优化这个地方。

流水线并行

解决时间问题

核心思想：在模型并行的基础上，引入数据并行。把原来的mini_batch再进行划分，称为micro_batch。

在这张图中，第一个下标表示GPU的编号，第二个下标表示micro_batch的编号。假设我们将mini_batch划分为了M个。那么在这张图下，整个流程的无效时间占比是：
只要M切片够大，那么整个流程的无效时间就会越来越少。

解决空间问题

前文提到，随着模型的增加，GPU存储的中间结果越来越大。对此，Gpipe采取了一种非常简单粗暴但有效的办法：用时间换空间，在论文里，这种方法被命名为re-materalization，后人也称其为active checkpoint.
具体来说，就是几乎不存中间结果，等到backward的时候，再重新算一遍forward，图例如下：

每块GPU上，我们只保存来自上一块的最后一层输入z，其余的中间结果我们算完就废。等到backward的时候再由保存下来的z重新进行forward来算出。
最终可以算出，每块GPU的峰值的空间复杂度为：
相比之下大大减少了对GPU内存的压力。在pytorch中，就是参数gradient_checkpoint所做的事情。

数据并行

流水线并行有一个问题，就是模型的均匀切割问题。模型能否被均匀的切割，决定了整体的计算效率。因此不够端到端。
基于此，发展出了数据并行的范式。
核心思想：在各个GPU上都拷贝一份完整模型，各自吃一份数据，算一份梯度，最后对梯度进行累加来更新整体模型
数据并行的不同实现方式：
- DP
- DDP
- ZeRO（现在主流）

DP实现

- 若干块计算GPU，如图中GPU0~GPU2；1块梯度收集GPU，如图中AllReduce操作所在GPU。
- 在每块计算GPU上都拷贝一份完整的模型参数。
- 把一份数据X（例如一个batch）均匀分给不同的计算GPU。
- 每块计算GPU做一轮FWD和BWD后，算得一份梯度G。
- 每块计算GPU将自己的梯度push给梯度收集GPU，做聚合操作。这里的聚合操作一般指梯度累加（一般是求平均）。当然也支持用户自定义。
- 梯度收集GPU聚合完毕后，计算GPU从它那pull下完整的梯度结果，用于更新模型参数W。更新完毕后，计算GPU上的模型参数依然保持一致。
- 聚合再下发梯度的操作，称为AllReduce。
实现DP操作的经典编程框架称为“参数服务器”，在这个框架中，计算GPU称为Worker，梯度聚合GPU称为Server。在实际的应用中，一般可以选择一个Worker同时作为Server。
额外的说明：
- 一个Worker或者Server下可以不止一块GPU
- Server可以只做梯度聚合，也可以梯度聚合+全量参数更新一起做

DP框架带来了两个问题
- 存储开销大。每块GPU上都存了一份完整的模型，造成冗余。关于这一点的优化，我们将在后文ZeRO部分做讲解。
- 通讯开销大。Server需要和每一个Worker进行梯度传输。当Server和Worker不在一台机器上时，Server的带宽将会成为整个系统的计算效率瓶颈。（这部分可以用异步梯度更新来解决，具体是在等待Server聚合梯度的过程，Worker继续拿数据，用旧的W进行计算。至于继续计算几轮，可以用延迟步数这个超参数决定。其实思想和之前的micro batch一样）

DDP

受通讯负载不均的影响，DP一般用于单机多卡场景。因此，DDP作为一种更通用的解决方案出现了，既能多机，也能单机。DDP首先要解决的就是通讯问题：将Server上的通讯压力均衡转到各个Worker上。
前文我们说过，聚合梯度 + 下发梯度这一轮操作，称为AllReduce。接下来我们介绍目前最通用的AllReduce方法：Ring-AllReduce。它由百度最先提出，非常有效地解决了数据并行中通讯负载不均的问题，使得DDP得以实现。

Ring-AllReduce

如下图，假设有4块GPU，每块GPU上的数据也对应被切成4份。AllReduce的最终目标，就是让每块GPU上的数据都变成箭头右边汇总的样子。

Ring-AllReduce分为两大步骤实现这个目标：
- Reduce-Scatter: 定义环形GPU拓扑关系，每次将自己得到的所有梯度数据传递给相邻的GPU节点。使得最终只需要GPU_NUM次就可以得到最终的数据。
- All-Gather:利用环形拓扑关系，把拥有完整数据的节点的数据广播到所有的节点。
大概就是不断执行这个操作

最后把红色的部分进行广播

ZeRO

通讯负载不均的优化我们在DDP中解决过了，但还遗留了一个显存开销问题：数据并行中，每个GPU上都复制了一份完整模型，当模型变大时，很容易打爆GPU的显存，那要怎么办呢？
ZeRO被用来解决大模型训练中的显存开销问题。主要思想是用通讯换显存。

存储分类

存储主要分为两大块：Model States和Residual States
Model States指和模型本身息息相关的，必须存储的内容，具体包括：

optimizer states：Adam优化算法中的momentum和variance
gradients：模型梯度
parameters：模型参数W

Residual States指并非模型必须的，但在训练过程中会额外产生的内容，具体包括：

activation：激活值。在流水线并行中我们曾详细介绍过。在backward过程中使用链式法则计算梯度时会用到。有了它算梯度会更快，但它不是必须存储的，因为可以通过重新做Forward来算它。
temporary buffers: 临时存储。例如把梯度发送到某块GPU上做加总聚合时产生的存储。
unusable fragment memory：碎片化的存储空间。虽然总存储空间是够的，但是如果取不到连续的存储空间，相关的请求也会被fail掉。对这类空间浪费可以通过内存整理来解决。

精度混合训练

对于模型，我们肯定希望其参数越精准越好，也即我们用fp32（单精度浮点数，存储占4byte） 来表示参数W。但是在forward和backward的过程中，fp32的计算开销也是庞大的。那么能否在计算的过程中，引入fp16或bf16（半精度浮点数，存储占2byte），来减轻计算压力呢？于是，混合精度训练就产生了，它的步骤如下图：

- 存储一份fp32的parameter，momentum和variance（统称model states）
- 在forward开始之前，额外开辟一块存储空间，将fp32 parameter减半到fp16 parameter。
- 正常做forward和backward，在此之间产生的activation和gradients，都用fp16进行存储。
- 用fp16 gradients去更新fp32下的model states。
- 当模型收敛后，fp32的parameter就是最终的参数输出。

总存储大小

因为采用了Adam优化，所以才会出现momentum和variance，当然你也可以选择别的优化办法。因此这里为了更通用些，记模型必存的数据大小为。因此最终内存开销为：
另外，这里暂不将activation纳入统计范围，原因是：

activation不仅与模型参数相关，还与batch size相关
activation的存储不是必须的。存储activation只是为了在用链式法则做backward的过程中，计算梯度更快一些。但你永远可以通过只保留最初的输入X，重新做forward来得到每一层的activation（虽然实际中并不会这么极端）。
因为activation的这种灵活性，纳入它后不方便衡量系统性能随模型增大的真实变动情况。因此在这里不考虑它。

ZeRO-DP

知道了什么东西会占存储，以及它们占了多大的存储之后，我们就可以来谈如何优化存储了。
注意到，在整个训练中，有很多states并不会每时每刻都用到，举例来说；

Adam优化下的optimizer states只在最终做update时才用到
数据并行中，gradients只在最后做AllReduce和updates时才用到
参数W只在做forward和backward的那一刻才用到
诸如此类

所以，ZeRO想了一个简单粗暴的办法：如果数据算完即废，等需要的时候，我再想办法从个什么地方拿回来，那不就省了一笔存储空间吗？
沿着这个思路，我们逐一来看ZeRO是如何递进做存储优化的。

$：优化状态分割$

我们知道，对于混合精度的训练过程。更新模型的时候需要用到：fp32的那些必存项和梯度信息，但是不更新模型的时候这些都不用。
因此，状态分割的核心思想就是，把最大头的所有的fp32的这部分拆开，拆到每一个GPU上。然后只对GPU上存有fp32信息的部分进行更新。
即：如下图所示

在进行完这个操作之后，我们对W进行All-Gather从别的GPU上把更新好的W取回来。

由此可以得到显存的情况如上表所示。

$优化状态与梯度分割$

其实很显然，我们说了，更新模型需要用到优化器状态和梯度状态，那你优化器都分割了，为啥不把梯度信息也分割了呢？

此时整体的流程如下：
(1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后，算得一份完整的梯度(注意，这里还没有分割，存的是完整的梯度)。
（2）对梯度做一次Reduce-Scatter，保证每个GPU上所维持的那块梯度是聚合梯度。例如对GPU1，它负责维护G1，因此其他的GPU只需要把G1对应位置的梯度发给GPU1做加总就可。汇总完毕后，白色块对GPU无用，可以从显存中移除
（3）每块GPU用自己对应的O和G去更新相应的W。更新完毕后，每块GPU维持了一块更新完毕的W。同理，对W做一次All-Gather，将别的GPU算好的W同步到自己这来。

不过这种优化方式，峰值的显存和上面一样，因为你还是有时刻是要储存一个完整的G信息的。

$：优化状态、梯度与参数分割$

那么循序渐进的，很容易我们可以想到，也可以把模型的fp16参数也给分割了
这里就不做过多的赘述了，应该很容易明白这部分的原理。

但是，我们需要明确的是，虽然切割了模型，但是该方法仍然不属于模型并行。
ZeRO是模型并行的形式，数据并行的实质。
模型并行，是指在forward和backward的过程中，我只需要用自己维护的那块W来计算就行。即同样的输入X，每块GPU上各算模型的一部分，最后通过某些方式聚合结果。
但对ZeRO来说，它做forward和backward的时候，是需要把各GPU上维护的W聚合起来的，即本质上还是用完整的W进行计算。它是不同的输入X，完整的参数W，最终再做聚合。

Reference

https://zhuanlan.zhihu.com/p/613196255 知乎猛猿

大模型并行技术

BIMU