在上一篇文章中，我们介绍了以Google GPipe为代表的流水线并行范式。当模型规模过大，单块GPU无法容纳时，流水线并行会将模型的不同层分配到不同GPU，通过切割mini-batch实现流水线处理，提升计算与通信的比率，同时借助re-materialization机制有效降低显存占用。
不过在实际应用中，流水线并行的使用并不十分广泛——核心原因在于模型能否均匀切割直接影响整体计算效率，这往往需要算法工程师手动调优。那么，有没有一种更简洁、更通用的并行方案呢？答案是肯定的。今天我们就来聊一聊应用最广泛、也最易于理解的一种并行范式：数据并行。
数据并行的核心思想非常直观：每块GPU上存放一份完整的模型副本，各自处理一份数据，计算出一份梯度，最后将梯度累加起来更新整体模型。理念并不复杂，但在大模型场景下，巨大的存储开销和GPU之间的通信量，却成为系统设计必须面对的硬骨头。下面我们递进式地介绍三种主流实现方式：
- DP（Data Parallelism） ：最早的数据并行模式，通常基于参数服务器（Parameter Server）编程框架，实际中多用于单机多卡场景。
- DDP（Distributed Data Parallelism） ：分布式数据并行，采用Ring-AllReduce通信方式，实际中多用于多机场景。
- ZeRO ：零冗余优化器，由微软推出并集成在DeepSpeed框架中。严格来说，ZeRO采用数据并行+张量并行的组合，旨在降低存储开销。
本文先讲解DP和DDP，ZeRO留到下一篇。全文结构如下：
1、数据并行（DP）
- 1.1 整体架构
- 1.2 通信瓶颈与梯度异步更新
2、分布式数据并行（DDP）
- 2.1 图解Ring-AllReduce
- 2.2 DP与DDP通信量分析
推荐阅读：
图解大模型训练之：流水线并行，以GPipe为例
一、数据并行（DP）
1.1 整体架构
一个经典的数据并行过程大致如下：
- 若干块计算GPU（例如图中的GPU0~GPU2）和1块梯度收集GPU（负责AllReduce操作）。
- 每块计算GPU上都拷贝一份完整的模型参数。
- 将一份数据X（比如一个batch）均匀分给不同的计算GPU。
- 每块计算GPU执行一次前向计算和反向传播后，得到一份梯度G。
- 每块计算GPU将自己的梯度push给梯度收集GPU，进行聚合操作。这里的聚合通常指梯度累加，当然也支持用户自定义。
- 梯度收集GPU聚合完毕后，计算GPU从它那里pull回完整的梯度结果，用于更新模型参数W。更新完成后，所有计算GPU上的模型参数保持一致。
- 聚合再下发梯度的这一整套操作，称为AllReduce。
实现DP的一种经典编程框架叫“参数服务器”：其中计算GPU称为Worker，梯度聚合GPU称为Server。实际中为了减少通信量，通常可以选一个Worker同时兼任Server——例如把梯度全发到GPU0上做聚合。这里有几个额外要点：
- 一个Worker或Server可以包含多块GPU。
- Server可以只做梯度聚合，也可以同时承担梯度聚合与全量参数更新的任务。
在参数服务器的语言体系下，DP的过程可以进一步用下图表示：
1.2 通信瓶颈与梯度异步更新
DP的框架理解起来并不难，但实际应用中会遭遇两个主要问题：
- 存储开销大：每块GPU上都需要存放一份完整模型，造成冗余。关于这一点，我们会在后面的ZeRO部分专门讨论。
- 通信开销大：Server需要与每一个Worker传输梯度。当Server和Worker不在同一台机器上时，Server的带宽会成为整个系统的性能瓶颈。
仔细想想通信开销这件事：如果把传输比作马路，带宽就是马路的宽度，决定了每次能并排通过多少数据。比如带宽是100G/s，但每秒却往Server塞1000G的数据，消化自然需要时间。
那么问题来了：能不能让Worker在等待梯度返回的这段时间里继续干活，而不是干等着？“梯度异步更新”正是为此而生。
上图展示了在梯度异步更新场景下，某个Worker的计算顺序：
- 在第10轮计算中，Worker正常计算梯度，并向Server发送push&pull请求。
- 但Worker并不会老老实实等到拿到聚合梯度、更新完参数W后再开始下一轮。它会直接拿着旧的W，吃新的数据，继续第11轮的计算。这样就能保证通信期间Worker也在马不停蹄地计算，从而提升计算与通信的比率。
- 当然，异步也不能太过分。只计算梯度不更新权重，模型根本无法收敛。图中展示的是延迟为1的异步更新——意味着在第12轮计算开始时，必须确保W已经用第10、11轮的梯度完成两次更新。
在参数服务器框架下，延迟的步数可以由用户自定义。下图分别刻画了几种延迟情况：
- (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来说：
- Reduce-Scatter阶段，通信量为
- All-Gather阶段，通信量为
总通信量为，随着N增大，可以近似为。
而前文DP的Server承载的总通信量也是。虽然通信量相同，但搬运相同数据量的时间却不一定相同。DDP将通信量平均分摊到了每一时刻的每个Worker上，而DP只让Server一个人当搬运工。当越来越多的GPU分布在相距较远的机器上时，DP的通信时间会显著增加。
这并不意味着参数服务器已经过时了（不少文章把参数服务器视为老古董）。实际上，参数服务器也提供了多Server方案，如下图所示：
在多Server模式下，每个Server可以只负责维护和更新某一块梯度（也可以连带维护该块参数）。虽然每个Server依然需要与所有Worker通信，但它的带宽压力会小很多。经过精心设计后，依然可以用于DDP。本文采用递进式的方式介绍DP和DDP，并不想决出孰优孰劣——方法本来就是多样性的。对参数服务器感兴趣的朋友，可以翻阅文末参考链接中的第一个。
最后，请记住Ring-AllReduce这个方法，因为在之后的ZeRO、Megatron-LM中，它会频繁出现，是分布式训练系统中一个重量级算子。
三、总结
1、在DP中，每个GPU上都存一份完整的模型副本，每个GPU处理batch的一部分数据，所有GPU算出的梯度累加后再传回各GPU用于更新参数。
2、DP多采用参数服务器编程框架，通常由若干计算Worker和一个梯度聚合Server组成。Server与每个Worker通信，Worker之间不直接通信。因此Server承担了系统全部的通信压力，这也决定了DP多用于单机多卡场景。
3、梯度异步更新是提升计算通信比的一种方法，延迟步数的大小会直接影响模型的收敛速度。
4、Ring-AllReduce通过构建网络环拓扑，将通信压力均匀地分摊到每个GPU上，使得跨机器的数据并行（DDP）得以高效实现。
5、DP和DDP的总通信量相同，但通信负载不均导致DP需要耗费更多时间来搬运数据。