专栏名称: 极市平台

极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台，为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯，行业动态，在线分享信息，线下活动等。网站: http://cvmart.net/

图解大模型训练之：数据并行（DP、DDP、ZeRO、零冗余优化)

极市平台 · 公众号 · · 2024-09-01 22:00

主要观点总结

本文介绍了大模型场景里数据并行的实现方法，包括DP、DDP和ZeRO三种方式。文章详细阐述了每种方式的原理、实现细节和优缺点，以及针对存储消耗的优化方法。最后简单介绍了ZeRO-Offload与ZeRO-Infinity。

关键观点总结

关键观点1: 数据并行的实现方式

包括DP、DDP和ZeRO三种。DP是最早的数据并行模式，采用参数服务器这一编程框架；DDP是更通用的解决方案，通过Ring-AllReduce方法解决通讯问题；ZeRO用通讯换显存，通过优化状态分割、梯度分割和参数分割来降低存储消耗。

关键观点2: 存储消耗的优化

包括优化状态分割、优化状态与梯度分割、优化状态、梯度与参数分割等。使用固定大小的内存buffer进行存储，提升带宽利用率，并设置机制对碎片化的存储空间进行重新整合。

关键观点3: ZeRO-Offload与ZeRO-Infinity

ZeRO-Offload将部分数据卸载到CPU上，以减少显存和通讯压力。ZeRO-infinity同理，它们都在解决数据存储和计算性能的问题。

正文

↑ 点击蓝字关注极市平台

作者丨猛猿

来源丨大猿搬砖简记

编辑丨极市平台

极市导读

大模型场景里巨大的存储和GPU间通讯量是系统设计时需要考虑的重点，本文递进介绍了三种主流数据并行的实现方法：DP、DD皮、ZeRo。 >> 加入极市CV技术交流群，走在计算机视觉的最前沿

当模型太大，一块GPU放不下时，流水线并行将模型的不同层放到不同的GPU上，通过切割mini-batch实现对训练数据的流水线处理，提升GPU计算通讯比。同时通过re-materialization机制降低显存消耗。

但在实际应用中，流水线并行并不特别流行，主要原因是模型能否均匀切割，影响了整体计算效率，这就需要算法工程师做手调。 因此，今天我们来介绍一种应用最广泛，最易于理解的并行范式：数据并行。

数据并行的核心思想是： 在各个GPU上都拷贝一份完整模型，各自吃一份数据，算一份梯度，最后对梯度进行累加来更新整体模型 。理念不复杂，但到了大模型场景， 巨大的存储和GPU间的通讯量， 就是系统设计要考虑的重点了。在本文中，我们将递进介绍三种主流数据并行的实现方式：

DP（Data Parallelism） ：最早的数据并行模式，一般采用参数服务器(Parameters Server)这一编程框架。实际中多用于单机多卡
DDP（Distributed Data Parallelism） ：分布式数据并行，采用Ring AllReduce的通讯方式，实际中多用于多机场景
ZeRO： 零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式，旨在降低存储。

一、数据并行（DP）

1.1 整体架构

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

若干块 计算GPU ，如图中GPU0~GPU2；1块 梯度收集GPU ，如图中AllReduce操作所在GPU。
在每块计算GPU上都拷贝一份完整的模型参数。
把一份数据X（例如一个batch）均匀分给不同的计算GPU。
每块计算GPU做一轮FWD和BWD后，算得一份梯度G。
每块计算GPU将自己的梯度 push 给梯度收集GPU，做聚合操作。这里的聚合操作一般指 梯度累加 。当然也支持用户自定义。
梯度收集GPU聚合完毕后，计算GPU从它那 pull 下完整的梯度结果，用于更新模型参数W。更新完毕后，计算GPU上的模型参数依然保持一致。
聚合再下发梯度的操作，称为AllReduce 。

前文说过，实现DP的一种经典编程框架叫“参数服务器”，在这个框架里， 计算GPU称为Worker ， 梯度聚合GPU称为Server。 在实际应用中，为了尽量减少通讯量，一般可选择一个Worker同时作为Server。比如可把梯度全发到GPU0上做聚合。需要再额外说明几点：

1个Worker或者Server下可以不止1块GPU。
Server可以只做梯度聚合，也可以梯度聚合+全量参数更新一起做

在参数服务器的语言体系下，DP的过程又可以被描述下图：

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

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

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

我们对通讯开销再做详细说明。如果将传输比作一条马路，带宽就是马路的宽度，它决定每次并排行驶的数据量。例如带宽是100G/s，但每秒却推给Server 1000G的数据，消化肯定需要时间。那么当Server在搬运数据，计算梯度的时候，Worker们在干嘛呢？当然是在：

人类老板不愿意了：“打工系统里不允许有串行存在的任务！”，于是 梯度异步更新 这一管理层略诞生了。

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

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

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

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

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

batch就像活，异步就像画出去的饼，且往往不指定延迟步数，每个Worker干越来越多的活，但模型却没收敛取效，这又是刺伤了哪些打工仔们的心（狗头

二、分布式数据并行(DDP)

受通讯负载不均的影响， DP一般用于单机多卡场景 。因此，DDP作为一种更通用的解决方案出现了，既能多机，也能单机。 DDP首先要解决的就是通讯问题：将Server上的通讯压力均衡转到各个Worker上。实现这一点后，可以进一步去Server，留Worker。

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

2.1 Ring-AllReduce

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

Ring-ALLReduce则分两大步骤实现该目标： Reduce-Scatter 和 All-Gather。

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个数为。则梯度大小也为 , 每个梯度块的大小为对单卡GPU来说（只算其send通讯量）：

Reduce-Scatter阶段, 通讯量为
All-Gather阶段, 通讯量为

单卡总通讯量为 , 随着的增大, 可以近似为。全卡总通讯量为

而对前文的DP来说, 它的Server承载的通讯量是 , Workers为 , 全卡总通讯量依然为。虽然通讯量相同, 但搬运相同数据量的时间却不一定相同。DDP把通讯量均衡负载到了每一时刻的每个Worker上, 而DP仅让Server做勤劳的搬运工。当越来越多的GPU分布在距离较远的机器上时, DP的通讯时间是会增加的。

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

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

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

三、总结

1、在DP中，每个GPU上都拷贝一份完整的模型，每个GPU上处理batch的一部分数据，所有GPU算出来的梯度进行累加后，再传回各GPU用于更新参数

2、DP多采用参数服务器这一编程框架，一般由若个计算Worker和1个梯度聚合Server组成。Server与每个Worker通讯，Worker间并不通讯。因此Server承担了系统所有的通讯压力。基于此DP常用于单机多卡场景。

3、异步梯度更新是提升计算通讯比的一种方法，延迟更新的步数大小决定了模型的收敛速度。

4、Ring-AllReduce通过定义网络环拓扑的方式，将通讯压力均衡地分到每个GPU上，使得跨机器的数据并行（DDP）得以高效实现。

5、DP和DDP的总通讯量相同，但因负载不均的原因，DP需要耗费更多的时间搬运数据

由微软开发的 ZeRO（零冗余优化） ，它是DeepSpeed这一分布式训练框架的核心，被用来解决大模型训练中的显存开销问题。 ZeRO的思想就是用通讯换显存。 如果初读ZeRO，觉得它逻辑跳跃，晦涩难懂，那么下文或许可以帮到你～

四、存储消耗

4.1 存储分类

首先，我们来看在大模型训练的过程中，GPU都需要存什么内容。

存储主要分为两大块：Model States和Residual States
Model States 指和模型本身息息相关的，必须存储的内容，具体包括：

optimizer states ：Adam优化算法中的momentum和variance
gradients ：模型梯度
parameters ：模型参数W

Residual States 指并非模型必须的，但在训练过程中会额外产生的内容，具体包括：

activation ：激活值。在流水线并行中我们曾详细介绍过。在backward过程中使用链式法则计算梯度时会用到。有了它算梯度会更快，但它不是必须存储的，因为可以通过重新做Forward来算它。
temporary buffers: 临时存储。例如把梯度发送到某块GPU上做加总聚合时产生的存储。
unusable fragment memory ：碎片化的存储空间。虽然总存储空间是够的，但是如果取不到连续的存储空间，相关的请求也会被fail掉。对这类空间浪费可以通过内存整理来解决。

4.2 精度混合训练

知道了存储分类，进一步，我们想知道，假设模型的参数W大小是，那么每一类存储具体占了多大的空间呢？

在分析这个问题前，我们需要来了解 精度混合训练 。

对于模型，我们肯定希望其参数越精准越好，也即我们用 fp32（单精度浮点数，存储占4byte） 来表示参数W。但是在forward和backward的过程中，fp32的计算开销也是庞大的。那么能否在计算的过程中，引入 fp16或bf16（半精度浮点数，存储占2byte） ，来减轻计算压力呢？于是，混合精度训练就产生了，它的步骤如下图：

存储一份fp32的parameter，momentum和variance（统称model states）
在forward开始之前，额外开辟一块存储空间，将fp32 parameter减半到fp16 parameter。
正常做forward和backward，在此之间产生的activation和gradients，都用fp16进行存储。
用fp16 gradients去更新fp32下的model states。
当模型收敛后，fp32的parameter就是最终的参数输出。

通过这种方式，混合精度训练在计算开销和模型精度上做了权衡。如果不了解fp32，fp16和bf16的细节也没关系，不影响下文的阅读。只要记住它们所占的存储空间和精度表达上的差异即可。

4.3 存储大小

现在，我们可以来计算模型在训练时需要的存储大小了，假设模型的参数W大小是， 以byte为单位 ，存储如下：

因为采用了Adam优化, 所以才会出现momentum和variance, 当然你也可以选择别的优化办法。因此这里为了更通用些, 记模型必存的数据大小为。因此最终内存开销为:

另外， 这里暂不将activation纳入统计范围 ，原因是：

activation不仅与模型参数相关，还与batch size相关
activation的存储不是必须的。存储activation只是为了在用链式法则做backward的过程中，计算梯度更快一些。但你永远可以通过只保留最初的输入X，重新做forward来得到每一层的activation（虽然实际中并不会这么极端）。
因为activation的这种灵活性，纳入它后不方便衡量系统性能随模型增大的真实变动情况。因此在这里不考虑它，在后面会单开一块说明对activation的优化。

五、ZeRO-DP

知道了什么东西会占存储，以及它们占了多大的存储之后，我们就可以来谈如何优化存储了。

注意到，在整个训练中，有很多states并不会每时每刻都用到，举例来说；

Adam优化下的optimizer states只在最终做update时才用到
数据并行中，gradients只在最后做AllReduce和updates时才用到
参数W只在做forward和backward的那一刻才用到
诸如此类

所以，ZeRO想了一个简单粗暴的办法： 如果数据算完即废，等需要的时候，我再想办法从个什么地方拿回来，那不就省了一笔存储空间吗？

沿着这个思路，我们逐一来看ZeRO是如何递进做存储优化的。

5.1 ：优化状态分割

首先，从 optimizer state开始优化。将optimizer state分成若干份，每块GPU上各自维护一份。这样就减少了相当一部分的显存开销。如下图：

复习一下，此时W=fp16，G=fp16，O=fp32。此时，整体数据并行的流程如下：

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后，各得一份梯度。

（2）对梯度做一次 AllReduce ， 得到完整的梯度G ，产生单卡通讯量。为了表达简明，这里通讯量我们就不再换算成byte了，而直接根据参数量来计算。对 AllReduce（reduce-scatter + all-gather） 。

（3）得到完整梯度G，就可以对W做更新。我们知道W的更新由optimizer states和梯度共同决定。 由于每块GPU上只保管部分optimizer states，因此只能将相应的W（蓝色部分）进行更新 。（2）和（3）可以用下图表示：

（4）此时，每块GPU上都有部分W没有完成更新（图中白色部分）。所以我们需要对W做一次 All-Gather ，从别的GPU上把更新好的部分W取回来。产生单卡通讯量。

做完后, 设GPU个数为，显存和通讯量的情况如下：

假设各变量大小如表格第二列所示，那么在增加1.5倍单卡通讯开销的基础上，将单卡存储降低了4倍。看起来是个还不错的trade-off，那么还能做得更好吗

5.2 ：优化状态与梯度分割

现在，更近一步，我们把梯度也拆开，每个GPU格子维护一块梯度。

此时，数据并行的整体流程如下：

（1）每块GPU上存一份完整的参数W。将一个batch的数据分成3份，每块GPU各吃一份，做完一轮foward和backward后， 算得一份完整的梯度（下图中绿色+白色） 。

（2）对梯度做一次 Reduce-Scatter ，保证每个GPU上所维持的那块梯度是聚合梯度。例如对GPU1，它负责维护G1，因此其他的GPU只需要把G1对应位置的梯度发给GPU1做加总就可。汇总完毕后，白色块对GPU无用，可以从显存中移除。单卡通讯量。（1）和（2）见下图：

（3）每块GPU用自己对应的O和G去更新相应的W。更新完毕后， 每块GPU维持了一块更新完毕的W 。同理，对W做一次 All-Gather ，将别的GPU算好的W同步到自己这来。单卡通讯量。

再次比对下显存和通讯量：

和朴素DP相比， 存储降了8倍，单卡通讯量持平 ，好像更牛皮了呢！那么，还可以优化吗？