数据并行的核心思想是:在各个GPU上都复制一份完整的模型,每个GPU处理一部分数据,计算一份梯度,最后把梯度加总来更新整体模型。这个概念很简单,但是对大型模型来说,巨大的存储空间和GPU之间的通信量就是系统设计要考虑的重点。本文将逐步介绍三种主流的数据并行实现方法:

•DP(数据并行):最早的数据并行模式,通常使用参数服务器框架。主要用于单机多卡场景。

•DDP(分布式数据并行):采用Ring AllReduce通信方式,主要用于多机场景。

•ZeRO:由微软开发,用于其DeepSpeed框架。严格来说,ZeRO采用数据并行+张量并行方法,目的是降低存储需求。

1. DP 

1.1 总体架构一个典型的数据并行过程如下:

•多个计算GPU,如图中的GPU0~GPU2;1个梯度聚合GPU,如图中的AllReduce操作所在GPU。

•在每个计算GPU上都复制一份完整的模型参数。 

•将一份数据X(例如一个batch)均匀分给不同的计算GPU。

•每个计算GPU进行一轮前向和反向传播后,计算一份梯度G。

•每个计算GPU将自己的梯度推送给梯度聚合GPU,进行聚合操作。这里的聚合操作通常是梯度加和。当然,也支持自定义操作。

•梯度聚合GPU聚合完成后,计算GPU从中拉取完整的梯度结果,用于更新模型参数W。更新后,计算GPU上的模型参数仍然保持一致。

 •聚合然后广播梯度的操作称为AllReduce。

如前所述,实现DP的一种典型编程框架是“参数服务器”。在这个框架中,计算GPU称为Worker,梯度聚合GPU称为Server。在实际应用中,为了最大限度地减少通信量,通常选择一个Worker同时作为Server。例如,可以把所有梯度发送到GPU0进行聚合。需要额外说明几点:

•一个Worker或Server下可以有不止一块GPU。

 •Server可以只进行梯度聚合,也可以同时进行梯度聚合和完整的参数更新。

1.2 缺点 

•存储开销大。每个GPU上都存有一份完整的模型,导致冗余。关于这一点的优化,我们将在后续的ZeRO部分进行讲解。 PS模式下的DP会造成负载不平衡,因为充当server的GPU需要一定的显存来保存worker节点计算出的局部梯度。

 •通信开销大。Server还需要将更新后的模型参数广播到每个Worker,Server的带宽就成了Server与Worker之间通信的瓶颈,Server与Worker之间的通信成本会随着Worker数量的增加线性增加。

 举个例子:假设我们有一个300M的参数模型,每个参数占4个字节,那么该模型占用内存为300M*4=1.2GB。假设网络带宽为1GB/s(10GbE),当使用单机双卡进行训练时,同步参数需要1.2s。当我们使用单机5卡进行训练时,同步参数需要1.2s*(5-1)=4.8s。这还仅仅是同步参数所需的时间,还需要考虑Server接收Worker梯度所需的时间。在这种情况下,Server与Worker之间的通信时间大约为10s。而且,在Server与Worker之间进行通信时,Worker无法进行计算,导致大量计算资源浪费。

2. DDP

由于通信负载不平衡的影响,DP通常用于单机多卡场景。因此,DDP作为一种更通的解决方案出现了,既能用于多机,也能用于单机。DDP首先要解决的就是通信问题:将Server上的通信压力均衡转移到各个Worker上。实现这一点后,可以进一步去除Server,只保留Worker。前面我们说过,聚合梯度+广播梯度的这一轮操作称为AllReduce。接下来我们介绍目前最通用的AllReduce方法:Ring-AllReduce。它由百度最先提出,非常有效地解决了数据并行中通信负载不平衡的问题,使DDP得以实现。

2.1 Ring-AllReduce如下图,假设有4块GPU,每块GPU上的数据(事实上是计算好的梯度)也对应被切成4份。AllReduce的最终目标,就是让每块GPU上的数据都变成箭头右边汇总的样子。

•Reduce-Scatter  定义网络拓扑关系,使得每个GPU只与相邻的两块GPU通信。每次发送对应位置的数据进行累加。每一次累加更新都形成一个拓扑环,因此被称为Ring。见到这可能会感到困惑,我们用图例把详细步骤画出来。 一次累加完成后,蓝色位置的数据块被更新,被更新的数据块将成为下一次更新的起点,继续执行累加操作。 3次更新后,每块GPU上都有一块数据获得了对应位置完整的聚合(图中红色)。此时,Reduce-Scatter阶段结束。进入All-Gather阶段。目标是把红色块的数据广播到其他GPU对应的位置上。

•All-Gather如名字中Gather所述,此操作仍采用“相邻GPU对应位置进行通信”的原则,但对应位置的数据不再执行相加,而是直接替换。All-Gather以红色块作为起点。 以此类推,同样经过3轮迭代后,使得每块GPU上都汇总到了完整的数据.

2.2 Ring-AllReduce的通信量分析 假设模型参数W的大小为Φ,GPU数量为N。则梯度大小也为Φ,每个梯度块的大小为Φ/N。对单卡GPU来说(只计算其发送通信量):

•Reduce-Scatter阶段,通信量为(N-1)Φ/N 

•All-Gather阶段,通信量为(N-1)Φ/N 

 •单卡总通信量为2(N-1)Φ/N,随着N的增加,可以近似为2Φ。全卡总通信量为2NΦ •而对前文的DP来说,其Server承担的通信量是NΦ,Workers为NΦ,全卡总通信量同样为2NΦ。虽然通信量相同,但搬运相同数据量的时间却不一定相同。DDP将通信量均衡负载到每一时刻的每个Worker上,而DP仅让Server充当勤劳的搬运工。当越来越多的GPU 分布在距离较远的机器上时,DP的通信时间会增加。

2.3 暂时总结

1、在DP中,每个GPU上都复制一份完整的模型,每个GPU上处理batch的一部分数据,所有GPU计算出的梯度进行加和后,再传回各GPU用于更新参数。  2、DP主要采用参数服务器框架,通常由多个计算Worker和1个梯度聚合Server组成。Server与每个Worker通信,Worker之间不通信。因此,Server承担了系统的全部通信压力。基于此,DP通常用于单机多卡场景。3、Ring-AllReduce通过定义网络环拓扑的方式,将通信压力均衡地分配到每个GPU上,使跨机器的数据并行(DDP)得以高效实现。 4、DP和DDP的总通信量相同,但由于负载不均衡,DP需要花费更多时间搬运数据。

3. ZeRO

ZeRO是由微软在其DeepSpeed框架中提出的一种优化数据并行的方法。ZeRO代表“zero redundancy optimizer”(零冗余优化器)。它通过将模型参数划分为若干部分,并只在某些GPU上存储这些部分,从而最大限度地减少参数冗余,达到优化存储的目的。 ZeRO将模型的参数分割成若干分区,并只在某些GPU上存储这些分区的参数。所有的GPU仍然会计算全局梯度,但每个GPU只需要存储与其相关的局部参数。梯度会在所有的GPU之间交换,以更新全局模型。这种方法可以最大限度地减少冗余的参数存储,从而节省显存。ZeRO在实现上主要采用了三种方法:

•ZeRO-Offload:将一些参数分区 offload 到 CPU 内存中存储。这可以减少 GPU 内存使用,但会增加 CPU-GPU 通信负载。

•ZeRO-Partitioned:将参数分区分配到不同的 GPU 上存储。每个 GPU 只存储一部分参数分区。这可以充分利用多个 GPU 的存储资源,减少参数冗余。

•ZeRO-Stage:将参数分为热参数分区和冷参数分区。热参数分区频繁访问,存储在 GPU 上;冷参数分区不频繁访问,offload 到系统内存存储。这可以根据参数的访问频率选择合适的存储位置, balances GPU 内存使用和 CPU-GPU 通信负载。ZeRO 通过上述三种方法,可以将数据并行训练中的参数冗余显著减小,从而可以应用到更大规模的模型。但ZeRO也增加了实现的复杂度,需要慎重权衡。