大模型训练中的数据并行技术详解：DP、DDP与ZeRO

大模型训练中的数据并行技术详解：DP、DDP与ZeRO

引用

CSDN

1.

https://blog.csdn.net/qq_41554005/article/details/145101652

在大模型训练中，数据并行（DP）、分布式数据并行（DDP）和零冗余优化器（ZeRO）是三种主流的数据并行实现方式。本文将详细介绍这三种方法的核心原理、实现方式以及实际应用中的优化策略。

英文缩写

• Data Parallelism（DP）数据并行
• Distributed Data Parallelism（DDP）分布式数据并行
• Zero Redundancy Optimizer（ZeRO）零冗余优化器
• FWD：前传
• BWD：后传

背景

• PP并行的原理：当模型太大，一块GPU放不下时，流水线并行将模型的不同执行阶段（一般为不同层）放到不同的GPU上，通过将mini-batch切割成更细粒度的micro-batch，实现对训练数据的流水线处理，提升GPU计算通讯比。同时通过re-materialization机制降低显存消耗。
• DP并行的优势（更强易用性应用更加广泛）：在实际应用中，流水线并行并不特别流行，主要原因是模型能否均匀切割，这影响了整体计算效率，也需要算法工程师做手调。因此，来介绍一种应用最广泛，最易于理解的并行范式：数据并行。
• 数据并行的核心思想是：在各个GPU上都拷贝一份完整模型，各自吃一份数据，算一份梯度，最后对梯度进行累加来更新整体模型。理念不复杂，但到了大模型场景，巨大的存储和GPU间的通讯量，就是系统设计要考虑的重点了。在本文以及后续文章中，我们将递进介绍三种主流数据并行的实现方式：
• DP：最早的数据并行模式，一般采用参数服务器这一编程框架。实际中多用于单机多卡
• DDP：分布式数据并行，采用Ring AllReduce的通讯方式，实际中多用于多机场景
• ZeRO：零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式，旨在降低存储

1 数据并行（DP）

1.1 整体架构

一个经典数据并行的过程如下：

• 若干块计算GPU（Worker）：如图中GPU0~GPU2
• 1块梯度收集GPU（Server）：如图中AllReduce操作所在GPU
• 在每块计算GPU上都拷贝一份完整的模型参数W WW
• 把一份数据X XX（例如一个batch）均匀分给不同的计算GPU：如X 0 ∼ X 2 X0\sim X2X0∼X2
• 每块计算GPU做一轮FWD和BWD后，算得一份梯度G GG，如G 0 ∼ 2 G0\sim2G0∼2
• 每块计算GPU将自己的梯度push给梯度收集GPU，做聚合操作。这里的聚合操作一般指梯度累加。当然也支持用户自定义
• 梯度收集GPU聚合完毕后，计算GPU从它那pull下完整的梯度结果，用于更新模型参数W WW。更新完毕后，计算GPU上的模型参数依然保持一致
• 聚合再下发梯度的操作，称为AllReduce。
  前文说过，实现DP的一种经典编程框架叫“参数服务器”，在这个框架里，计算GPU称为Worker，梯度聚合GPU称为Server。在实际应用中，为了尽量减少通讯量，一般可选择一个Worker同时作为Server。比如可把梯度全发到GPU0上做聚合。需要再额外说明几点：
• 1个Worker或者Server下可以不止1块GPU，可以是多个GPU的集群。
• Server可以只做梯度聚合，也可以梯度聚合+全量参数更新一起做
  在参数服务器的语言体系下，DP的过程又可以被描述下图：

1.2 通讯量分析

假设模型参数W WW的大小为Φ \PhiΦ，GPU个数为N NN。则梯度大小也为Φ \PhiΦ，每个梯度块的大小为Φ N \frac{\Phi}{N}NΦ

对单个GPU（Worker）来说：

• push的过程需要完成每一个GPU上完整梯度（大小为Φ \PhiΦ）上送
• pull的过程需要完成每一个GPU上完整新权重（大小为Φ \PhiΦ）下拉
• 单卡总通讯量为2 Φ 2\Phi2Φ（包含上送和下拉）
• 全Worker卡总通讯量为2 N Φ 2N\Phi2NΦ

对单个GPU（Server）来说：

• push的过程需要完成每一个Worker上完整梯度上送的收集（大小为N Φ N\PhiNΦ）
• pull的过程需要完成对每一个Worker上完整新权重的下发（大小为N Φ N\PhiNΦ）
• 单卡总通讯量为2 N Φ 2N\Phi2NΦ

1.3 通讯瓶颈与梯度异步更新

DP的框架理解起来不难，但实战中确有两个主要问题：

• 存储开销大：每块GPU上都存了一份完整的模型，造成冗余。关于这一点的优化，在后文ZeRO部分做讲解
• 通讯开销大：Server需要和每一个Worker进行梯度传输。当Server和Worker不在一台机器上时，Server的带宽将会成为整个系统的计算效率瓶颈

对通讯开销再做详细说明。如果将传输比作一条马路，带宽就是马路的宽度，它决定每次并排行驶的数据量。例如带宽是100GB/s，但每秒却推给Server 1000GB的数据，消化肯定需要时间10s。那么当Server在搬运数据，计算梯度的时候，Worker们就在休息。于是梯度异步更新诞生了。

上图刻画了在梯度异步更新的场景下，某个Worker的计算顺序为：

• 在第10轮计算中，该Worker正常计算梯度，并向Server发送push&pull梯度请求。
• 但是，该Worker并不会实际等到把聚合梯度拿回来，更新完参数获得第10轮更新结果W 10 W_{10}W10 后再做计算。而是直接拿旧的权重比如W 9 W_9W9 ，吃新的数据，一边穿一边继续第11轮的计算。这样就保证在通讯的时间里，Worker也在马不停蹄做计算，提升计算通讯比。
• 当然，异步也不能太过份。只计算梯度，不更新权重，那模型就无法收敛。图中刻画的是延迟为1的异步更新，也就是在开始第12轮的计算时，必须保证W WWW已经用第10轮的梯度做完更新了。

参数服务器的框架下，延迟的步数也可以由用户自己决定，下图分别刻划了几种延迟情况：

• (1) 无延迟
• (2) 延迟但不指定延迟步数。也即在迭代2时，用的可能是老权重，也可能是新权重，听天由命
• (3) 延迟且指定延迟步数为1。例如做迭代3时，可以不拿回迭代2的梯度，但必须保证迭代0、1的梯度都已拿回且用于参数更新

总结一下，异步很香，但对一个Worker来说，只是等于W WWW不变，batch的数量增加了而已，在SGD下，会减慢模型的整体收敛速度。异步的整体思想是，比起让Worker闲着，倒不如让它多吃点数据，虽然反馈延迟了，但只要它在干活在学习就行。

2 分布式数据并行（DDP）

受通讯负载不均的影响，DP一般用于单机多卡场景。因此，DDP作为一种更通用的解决方案出现了，既能多机，也能单机。

DDP首先要解决的就是通讯问题：将Server上的通讯压力均衡转到各个Worker上。实现这一点后，可以进一步去Server，留Worker。

前文我们说过，聚合梯度 + 下发梯度这一轮操作，称为AllReduce。接下来我们介绍目前最通用的AllReduce方法：Ring-AllReduce。它由百度最先提出，非常有效地解决了数据并行中通讯负载不均的问题，使得DDP得以实现。

2.1 Ring-AllReduce（reduce-scatter + all-gather）

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

Ring-ALLReduce则分两大步骤实现该目标：

• Reduce-Scatter
• All-Gather

2.1.1 Reduce-Scatter（散射规约）

定义网络拓扑关系，使得每个GPU只和其相邻的两块GPU通讯。每次发送对应位置的数据进行累加。每一次累加更新都形成一个拓扑环，因此被称为Ring。看到这觉得困惑不要紧，我们用图例把详细步骤画出来。

一次累加完毕后，蓝色位置的数据块被更新，被更新的数据块将成为下一次更新的起点，继续做累加操作。

3次更新之后，每块GPU上都有一块数据拥有了对应位置完整的聚合（图中红色）。此时，Reduce-Scatter阶段结束。进入All-Gather阶段。目标是把红色块的数据广播到其余GPU对应的位置上。

2.1.2 All-Gather

如名字里Gather所述的一样，这操作里依然按照“相邻GPU对应位置进行通讯”的原则，但对应位置数据不再做相加，而是直接替换。All-Gather以红色块作为起点。

以此类推，同样经过3轮迭代后，使得每块GPU上都汇总到了完整的数据，变成如下形式：

2.2 Ring-AllReduce通讯量分析

假设模型参数W WW的大小为Φ \PhiΦ，GPU个数为N NN。则梯度大小也为Φ \PhiΦ，每个梯度块的大小为Φ N \frac{\Phi}{N}NΦ （如前文的例子，分成了四份数据，单卡需要完成的数据散射和归并是模型任务量是1/4）。

对单个GPU来说：

• Reduce-Scatter阶段，通讯量为( N − 1 ) Φ N \frac{(N−1)\Phi}{N}N(N−1)Φ （每一个数据块大小Φ N \frac{\Phi}{N}NΦ ，ring上走N − 1 N-1N−1次就可以在某一个GPU上完成归并）（如果考虑发送与收到则需要两份如此通讯量）
• All-Gather阶段，通讯量为( N − 1 ) Φ N \frac{(N−1)\Phi}{N}N(N−1)Φ （每一个数据块大小Φ N \frac{\Phi}{N}NΦ ，ring上走N − 1 N-1N−1块数据次就可以把某一块GPU上完成归并的数据块下发至每一个GPU）（如果考虑发送与收到则需要两份如此通讯量）
• 单卡单向总通讯量为2 ( N − 1 ) Φ N \frac{2(N−1)\Phi}{N}N2(N−1)Φ ，随着N NN的增大，可以近似为2 Φ 2\Phi2Φ，双向则为4 Φ 4\Phi4Φ
• 全卡单向总通讯量为2 N Φ 2N\Phi2NΦ，双向则为4 N Φ 4N\Phi4NΦ

一般互联带宽为双向带宽，即同时实现相同带宽的收与发。假设收与发的带宽均为B BB。

DP 收 发 收发通讯时间
Server（单卡） N

        Φ 
       
      
     
       N\PhiNΦ  N 
       
      
        Φ 
       
      
     
       N\PhiNΦ  N 
         
        
          Φ 
         
        
       
         B 
        
       
      
     
       \frac{N\Phi}{B}BNΦ​  

Worker（单卡） Φ

       \PhiΦ  Φ 
       
      
     
       \PhiΦ  Φ 
        
       
         B 
        
       
      
     
       \frac{\Phi}{B}BΦ​  

集群视角（以server为瓶颈） N

        Φ 
       
      
     
       N\PhiNΦ  N 
       
      
        Φ 
       
      
     
       N\PhiNΦ  N 
         
        
          Φ 
         
        
       
         B 
        
       
      
     
       \frac{N\Phi}{B}BNΦ​  

DDP 收 发 收发通讯时间
单卡Reduce-Scatter阶段 Φ

       \PhiΦ  Φ 
       
      
     
       \PhiΦ  Φ 
        
       
         B 
        
       
      
     
       \frac{\Phi}{B}BΦ​  

单卡All-Gather阶段 Φ

       \PhiΦ  Φ 
       
      
     
       \PhiΦ  Φ 
        
       
         B 
        
       
      
     
       \frac{\Phi}{B}BΦ​  

单卡All-Reduce 2

        Φ 
       
      
     
       2\Phi2Φ  2 
       
      
        Φ 
       
      
     
       2\Phi2Φ  2 
         
        
          Φ 
         
        
       
         B 
        
       
      
     
       \frac{2\Phi}{B}B2Φ​  

集群视角All-Reduce阶段 2

        N 
       
      
        Φ 
       
      
     
       2N\Phi2NΦ  2 
       
      
        N 
       
      
        Φ 
       
      
     
       2N\Phi2NΦ  2 
         
        
          N 
         
        
          Φ 
         
        
        
        
          N 
         
        
          B 
         
        
       
      
        = 
       
       
        
        
          2 
         
        
          Φ 
         
        
       
         B 
        
       
      
     
       \frac{2N\Phi}{NB}=\frac{2\Phi}{B}NB2NΦ​=B2Φ​  

DP中Server为瓶颈，搬运数据量均阻塞在Server的通讯能力上。DDP把通讯量均衡负载到了每一时刻的每个Worker上，当越来越多的GPU分布在距离较远的机器上时，DP的通讯时间是会增加的，但是DDP可以基本不变。

但这并不说明参数服务器不能打（有很多文章将参数服务器当作old dinosaur来看）。事实上，参数服务器也提供了多Server方法，如下图：

在多Server的模式下，进一步，每个Server可以只负责维护和更新某一块梯度（也可以某块梯度+参数一起维护），此时虽然每个Server仍然需要和所有Worker通讯，但它的带宽压力会小非常多。经过调整设计后，依然可以用来做DDP。虽然这篇文章是用递进式的方式来介绍两者，但不代表两者间一定要决出优劣。我想表达的观点是，方法是多样性的。对参数服务器有兴趣的朋友，可以阅读参考的第1个链接。

最后，请大家记住Ring-AllReduce的方法，因为在之后的ZeRO，Megatron-LM中，它将频繁地出现，是分布式训练系统中重要的算子。

3 总结

• 在DP中，每个GPU上都拷贝一份完整的模型，每个GPU上处理batch的一部分数据，所有GPU算出来的梯度进行累加后，再传回各GPU用于更新参数
• DP多采用参数服务器这一编程框架，一般由若个计算Worker和1个梯度聚合Server组成。Server与每个Worker通讯，Worker间并不通讯。因此Server承担了系统所有的通讯压力。基于此DP常用于单机多卡场景。
• 异步梯度更新是提升计算通讯比的一种方法，延迟更新的步数大小决定了模型的收敛速度。
• Ring-AllReduce通过定义网络环拓扑的方式，将通讯压力均衡地分到每个GPU上，使得跨机器的数据并行（DDP）得以高效实现。
• DP和DDP的总通讯量相同，但因负载不均的原因，DP需要耗费更多的时间搬运数据

4 参考

• https://arxiv.org/pdf/1910.02054.pdf
• https://blog.51cto.com/u_14691718/5631471
• https://blog.csdn.net/qq_43307074/article/details/127688761
• https://web.eecs.umich.edu/~mosharaf/Readings/Parameter-Server.pdf
• https://zh.d2l.ai/chapter_computational-performance/parameterserver.html
• https://blog.csdn.net/dpppBR/article/details/80445569
• https://arxiv.org/abs/1910.02054
• https://blog.51cto.com/u_14691718/5631471