运用AI集群,使机器学习算法更好地从大数据中练习出功能优秀的大模型是散布式机器学习的方针。为了完成该方针,一般需求根据硬件资源与数据/模型规模的匹配情况,考虑对核算使命、练习数据和模型进行划分,进行散布式存储和散布式练习。下面将对现有的并行练习办法进行简要介绍。

数据并行

数据并行是最常见的并行方式,因为它很简单。在数据并行练习中,数据集被分割成几个碎片,每个碎片被分配到一个设备上。这相当于沿批次维度对练习过程进行并行化。每个设备将持有一个完好的模型副本,并在分配的数据集碎片上进行练习。在反向传播之后,模型的梯度将被全部削减,以便在不同设备上的模型参数能够坚持同步。

分布式训练并行技术

模型并行

在数据并行练习中,一个明显的特点是每个 GPU 持有整个模型权重的副本。这就带来了冗余问题。另一种并行模式是模型并行,即模型被分割并散布在一个设备阵列上。一般有两种类型的并行:张量并行和流水线并行。张量并行是在一个操作中进行并行核算,如矩阵-矩阵乘法。流水线并行是在各层之间进行并行核算。因此,从另一个角度来看,张量并行能够被看作是层内并行,流水线并行能够被看作是层间并行。

张量并行

张量并行练习是将一个张量沿特定维度分红N块,每个设备只持有整个张量的1/N,一起不影响核算图的正确性。这需求额定的通信来确保成果的正确性。

以一般的矩阵乘法为例,假定咱们有C = AB。咱们能够将B沿着列分割成[B0 B1 B2 … Bn],每个设备持有一列。然后咱们将A与每个设备上B中的每一列相乘,咱们将得到[AB0 AB1 AB2 … ABn]。此刻,每个设备仍然持有一部分的成果,例如,设备(rank=0)持有AB0。为了确保成果的正确性,咱们需求搜集全部的成果,并沿列维串联张量。通过这种方式,咱们能够将张量散布在设备上,一起确保核算流程坚持正确。

分布式训练并行技术

流水线并行

流水线并行的中心思维是,模型按层分割成若干块,每块都交给一个设备。

在前向传递过程中,每个设备将中心的激活传递给下一个阶段。

在后向传递过程中,每个设备将输入张量的梯度传回给前一个流水线阶段。这答应设备一起进行核算,并增加了练习的吞吐量。

分布式训练并行技术

流水线并行练习的一个缺点是练习设备简单呈现闲暇状况,导致核算资源的浪费,加快效率没有数据并行高。

分布式训练并行技术

优化器相关的并行

现在随着模型越来越大,单个GPU的显存现在一般无法装下那么大的模型了。那么就要想办法对占显存的地方进行优化。一般来说,模型练习的过程中,GPU上需求进行存储的参数包含了模型自身的参数,优化器状况,激活函数的输出值,梯度,以及一些零时的Buffer。数据的占比如下图所示:

分布式训练并行技术

能够看到模型参数仅占一部分的,当进行混合精度运算时,其中模型状况参数(优化器状况+ 梯度+ 模型参数)占到了一大半以上。因此,咱们需求想办法去除冗余数据。

优化器相关的并行便是一种去除冗余数据的并行方案。现在这种并行最流行的办法是ZeRO(即零冗余优化器)。针对模型状况的存储优化(去除冗余),ZeRO运用的办法是分片,即每张卡只存1/N的模型状况量,这样体系内只保护一份模型状况。ZeRO有三个不同等级,对模型状况进行不同程度的分片:

  • ZeRO-1 : Optimizer States Sharding
  • ZeRO-2 : Optimizer States & Gradients Sharding
  • ZeRO-3 : Optimizer States & Gradients & Parameters Sharding

分布式训练并行技术

异构体系并行

上述的办法中,一般需求很多的 GPU 来练习一个大型模型。然而,人们常常疏忽的是,与 GPU 相比,CPU 的内存要大得多。在一个典型的服务器上,CPU 能够轻松具有几百GB的内存,而每个 GPU 一般只要16或32GB的内存。这促进人们考虑为什么 CPU 内存没有被用于散布式练习。

最近的发展是依托 CPU 乃至是 NVMe 磁盘来练习大型模型。首要的主意是,在不运用张量时,将其卸载回 CPU 内存或 NVMe 磁盘。通过运用异构体系架构,有可能在一台机器上容纳一个巨大的模型。

分布式训练并行技术

多维混合并行

多维混合并行指将数据并行、模型并行和流水线并行结合起来进行散布式练习。

分布式训练并行技术

总之,为了充分运用带宽,一般情况下,不同机架之间是数据并行(每个机架都有完好的模型),机架内部的服务器之间是流水线并行。同一服务器的不同加快卡之间是算子并行。这样做的原因是,算子级并行所需的通信量最大,流水线并行次之,数据并行所需通信量最小。

分布式训练并行技术