大模型训练方法:

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),即该专家在一个批次上的门控输出值的和。并且定义损失项加入到模型的总损失上。该损失项等于所有专家重要度的方差的平方,再加上一个手工调节的比例因子。这个损失项会鼓励所有专家有相同的重要度。

大模型python 大模型技术_数据

2. Megatron-LM

Training Multi-Billion Parameter Language Models Using Model Parallelism 2019

motivation

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

method

矩阵分解

大模型python 大模型技术_深度学习_02

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只需要更新自己负责的部分参数。

大模型python 大模型技术_Network_03

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结构

大模型python 大模型技术_大模型python_04

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 的速度优势非常明显。
  • 大模型python 大模型技术_深度学习_05

  •  
  • 大模型python 大模型技术_可扩展性_06