大模型分布式训练技术_outrageously large neural networks: the sparsely-g

 大模型训练方法：

1. MoE

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.

motivation

现在的模型越来越大，训练样本越来越多，每个样本都需要经过模型的全部计算，这就导致了训练成本的平方级增长。为了解决这个问题，即将大模型拆分成多个小模型，对于一个样本来说，无需经过所有的小模型去计算，而只是激活一部分小模型进行计算，这样就节省了计算资源。那么如何决定一个样本去经过哪些小模型呢？这就引入了一个稀疏门机制，即样本输入给这个门，得到要激活的小模型索引，这个门需要确保稀疏性，从而保证计算能力的优化。

method

首先MoE是一个层，而不是一整个模型。其次，正如我们刚才所说，这个模型结构包含一个门网络来决定激活哪个expert，同时包含n个expert网络，这n个expert网络一般是同结构的。

保证稀疏性和均衡性，对这个softmax做了处理,第一个变动就是KeepTopK，这是个离散函数，将top-k之外的值强制设为负无穷大，从而softmax后的值为0。第二个变动是加了noise，这个的目的是为了做均衡，这里引入了一个Wnoise的参数，后面还会在损失函数层面进行改动。

如果不做改进，那么这么多的expert，只有几个expert会被集中使用。为了改进这一问题，采用软性约束方法。我们定义对于一个批次训练样本的专家重要度（the importance of an expert），即该专家在一个批次上的门控输出值的和。并且定义损失项加入到模型的总损失上。该损失项等于所有专家重要度的方差的平方，再加上一个手工调节的比例因子。这个损失项会鼓励所有专家有相同的重要度。

2. Megatron-LM

Training Multi-Billion Parameter Language Models Using Model Parallelism 2019

motivation

利用模型并行方法解决大模型训练的问题，侧重张量Tensor并行提升训练速度。实现了简单而高效的模型并行方法，成功突破了传统单个GPU训练的限制。

method

矩阵分解

3. ZeRO

Memory Optimizations Toward Training Trillion Parameter Models 2019

motivation

解决每个gpu都存整个模型的参数，导致显存超出的问题，减少显存，增加计算开销。当前分布式训练主要的三种并行模式：数据并行、模型并行和流水线并行。在三种并行方式中，数据并行因其易用性，得到了最为广泛的应用。然而，数据并行会产生大量冗余 Model States 的空间占用。

1. Optimizer States: Optimizer States 是 Optimizer 在进行梯度更新时所需要用到的数据，例如 SGD 中的Momentum以及使用混合精度训练时的Float32 Master Parameters。

2. Gradient：在反向传播后所产生的梯度信息，其决定了参数的更新方向。

3. Model Parameter: 模型参数，也就是我们在整个过程中通过数据“学习”的信息。在传统数据并行下，每个进程都使用同样参数来进行训练。每个进程也会持有对Optimizer States等完整拷贝，同样占用了大量显存。ZeRO 则在数据并行的基础上，引入了对冗余Model States的优化。使用 ZeRO 后，各个进程之后只保存完整状态的1/GPUs，互不重叠，不再存在冗余。

method

ZeRO 有三个不同级别，分别对应对 Model States 不同程度的分割 (Paritition)：

• ZeRO-1：分割Optimizer States；

• ZeRO-2：分割Optimizer States与Gradients；

• ZeRO-3：分割Optimizer States、Gradients与Parameters

最后我们总结zero优化后的训练步骤流程，我们以Pos+g+p这个模式为例：
初始状态:
假设有Nd个gpu，则每个gpu（gpu_n）保存总参数的1/Nd;并且保存这些参数对应的梯度、优化器状态(P_n,G_n,O_n)；参数可以按层划分；每个gpu_n同时还负责分配到自己身上的数据data_n（数据并行）；
a.正向计算第一层时，gpu_n将自己负责的参数(P_n)广播给其它所有的gpu；后面的模型层以此类推；最后每个gpu_n获得自己对应数据data_n的loss_n；
b.进行反向计算，此时需要gpu_n将自己负责的参数(P_n)广播给其它所有的gpu，最后计算得到对应于数据data_n的梯度；
c.将第二步的梯度聚合到对应的gpu_n上，每个gpu负责更新自己的P_n,G_n,O_n；
进行下一次迭代
从这个训练流程，可以看出，传递参数过程中它使用了广播的方式，而梯度聚合过程则使用了类似Allreduce中reduce-scatter的模式，但不再需要进行all-gather，因为每个gpu只需要更新自己负责的部分参数。

4. GShard

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding ICLR 2021

motivation

第一个把MoE结构引入Transformer结构的工作,GShard将Transformer中的Feedforward Network（FFN）层替换成了MoE层，并且将MoE层和数据并行巧妙地结合起来。在数据并行训练时，模型在训练集群中已经被复制了若干份。GShard通过将每路数据并行的FFN看成MoE中的一个专家来实现MoE层，这样的设计通过在多路数据并行中引入All-to-All通信来实现MoE的功能。

method

其实就是把原来的FFN（两层全连接）替换成了MoE结构

5. Switch Transformers

SWITCH TRANSFORMERS: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity 2021

motivation

将语言模型的参数量扩展至 1.6 万亿,简化了MoE路由算法,只将 token 表征发送给单个专家。研究表明，这种简化策略保持了模型质量，降低了路由计算，并且性能更好。将这种 k=1 的策略称为 Switch 层。如果将太多 token 发送给一个专家（下文称为「丢弃的 token」），则会跳过计算，token 表征通过残差连接直接传递到下层。但增加专家容量也不是没有缺点，数值太高将导致计算和内存浪费，最后应用到t5模型上。

method

• 简化稀疏路由：只将 token 表征发送给单个专家

• 高效稀疏路由：太多 token 发送给一个专家，则会跳过计算，token 表征通过残差连接直接传递到下层

• 提升训练和微调的技巧：正则化大型稀疏模型，对大型稀疏模型使用可选择行精度

• 预训练可扩展性：专家的数量是扩展模型最有效的维度

• 基于步数的可扩展性：在保持每个 token 的 FLOPS 不变时，拥有更多的参数（专家）可以提高训练速度。

• 基于时间的可扩展性:在训练时间和计算成本都固定的情况下，Switch Transformer 的速度优势非常明显。