Zen: Empowering Distributed Training with Sparsity-driven Data Synchronization

Zhuang Wang, Rice University; Zhaozhuo Xu, Stevens Institute of Technology; Jingyi Xi, unaffiliated; Yuke Wang, Anshumali Shrivastava, and T. S. Eugene Ng, Rice University

A1 主要贡献

论文针对分布式训练中的梯度同步问题,揭示了稀疏张量的高稀疏性特性,并探讨了其同步方案的设计空间。研究目标是桥接稀疏张量同步的性能差距,找出充分利用稀疏性的最优通信方案。具体创新点包括:分析流行模型中稀疏张量的特性,如重叠比率、致密化比率和非零梯度分布的偏斜性;系统探索同步方案的四个基本维度(通信、聚合、分区、平衡),证明存在通信最优方案;基于这些发现,开发ZEN系统,实现近最优通信时间,通过数据无关的分层散列算法和高效编码方案处理负载平衡和索引开销。ZEN在通信时间上比最先进方法加速高达5.09倍,在训练吞吐量上加速高达2.48倍。

Figure 1: ZEN系统概述。
Figure 1: ZEN系统概述。

表1:三种具有自然张量稀疏性的DL模型。

表1图片
表1图片

表2:三种LLM的配置。AH:注意力头。

表2图片
表2图片

A3 背景知识/关键Observation/设计原则

2.1 稀疏张量的特性

本节分析自然梯度计算和稀疏化算法中的稀疏张量特性。 对于自然张量稀疏性,我们研究表1中三种DL模型的嵌入层梯度张量。对于梯度压缩产生的稀疏性,我们研究表2中三种大型语言模型(LLM),并应用DGC【[39],Deep gradient compression: Reducing the communication bandwidth for distributed training,2017,The International Conference on Learning Representations (ICLR)】选择梯度张量的前5%梯度【[8],L-GreCo: Layerwise-adaptive gradient compression for efficient and accurate deep learning,2022,无URL;[27],Beyond throughput and compression ratios: Towards high end-to-end utility of gradient compression,2024,无URL】。详细工作负载见§5.1。

Figure 2: DL模型中稀疏张量的特性。(a) 显示稀疏张量的重叠比率变化;(b) 显示张量在聚合后密度更高。
Figure 2: DL模型中稀疏张量的特性。(a) 显示稀疏张量的重叠比率变化;(b) 显示张量在聚合后密度更高。

定义1(稠密张量)。 我们将DL层中的原始梯度张量定义为稠密张量。

我们将梯度张量的密度定义为其非零梯度值的百分比。 当许多参数具有零梯度时,我们可以用稀疏格式表示梯度张量。典型的稀疏格式是坐标列表(COO),它存储非零梯度的列表和相应索引的列表【[23],Efficient sparse collective communication and its application to accelerate distributed deep learning,2021,ACM SIGCOMM 2021 Conference;[77],GRACE: A compressed communication framework for distributed machine learning,2021,IEEE 39th International Conference on Distributed Computing Systems (ICDCS)】。

定义2(稀疏张量)。 我们将稀疏格式的梯度张量定义为稀疏张量。

我们用M表示稠密张量G的大小,其密度为$d_G$,训练涉及n个节点。 为简单起见,本节假设每个节点只有一个GPU。

C1:稀疏张量的重叠变化。 与稠密张量类似,稀疏张量在同步期间被聚合。当聚合稠密张量时,不同GPU的梯度索引相同。然而,由于不同GPU上的训练输入批次不同,稀疏张量中非零梯度的索引事先未知。它们可能有重叠,但重叠程度取决于许多因素,如DL模型、训练数据集和批次。我们定义重叠比率【[68],A survey on similarity measures in text mining,2016,Machine Learning and Applications: An International Journal】来量化两个稀疏张量之间的重叠。

定义3(重叠比率)。 给定两个稀疏张量及其非零梯度索引集分别为$I_1$和$I_2$,它们的重叠比率定义为$|I_1 \cap I_2| / \min\{|I_1|, |I_2|\}$,其中$| \cdot |$是集合的基数。

图2a显示了本文研究的六个DL模型的重叠比率的概率密度函数(PDF)。 一个模型中的重叠比率近似正态分布,并且范围很广。此外,不同模型有不同的重叠比率分布。

C2:聚合后的张量大小变化。 当聚合稠密张量时,聚合前后张量大小保持不变。然而,当聚合稀疏张量时,稀疏张量的未知重叠导致聚合后张量大小变化。因为聚合涉及多个GPU的稀疏张量,我们将n个GPU张量聚合后的密度表示为$d_n^G$。我们观察到稀疏张量在聚合后变得更稠密。我们定义致密化比率来量化这一特性。

Figure 3: 非零梯度的分布是偏斜的。(a) 非零梯度分布的热图;(b) 偏斜比率。
Figure 3: 非零梯度的分布是偏斜的。(a) 非零梯度分布的热图;(b) 偏斜比率。

定义4(致密化比率)。 给定稠密张量G,其致密化比率定义为$\gamma_n^G = d_n^G / d_G$。

图2b展示了本文研究的DL模型的平均致密化比率$\gamma_n^G$与GPU数量的关系。 致密化比率随着GPU数量增加而增加,表明张量在聚合后密度更高。我们还可以看到致密化比率小于GPU数量,即$\gamma_n^G < n$。这表明不同GPU稀疏张量中非零梯度的索引部分重叠。

C3:非零梯度的分布是偏斜的。 当均匀地将稠密张量分割成多个分区时,我们观察到大多数非零梯度都在其中一个分区中。例如,在八个分区中,六个DL模型中第一个分区中有超过60%的非零梯度。图3a显示了每个分区中非零梯度的百分比热图。我们定义偏斜比率来量化非零梯度的偏斜分布。

Figure 4: 四个GPU的三种通信模式的说明。GPU P3聚合所有GPU的数据。
Figure 4: 四个GPU的三种通信模式的说明。GPU P3聚合所有GPU的数据。

定义5(偏斜比率)。 给定稠密张量G,我们均匀地将G分成n个不相交的分区,表示为$\{G_1, \cdots, G_n\}$,G的n分区偏斜比率定义为$\max_{i \in [n]} \{ d_{G_i} \} / d_G$。

图3b展示了本文研究的DL模型中梯度张量的偏斜比率。 它们在所有六个模型中都很显著。例如,当我们将LSTM嵌入表中的梯度张量均匀分割成128个分区时,偏斜比率超过70。这表明超过一半的非零梯度在同一个分区中。

2.2 同步的基本维度

稠密张量的同步已被广泛研究【[29],A unified architecture for accelerating distributed DNN training in heterogeneous GPU/CPU clusters,2020,14th USENIX Symposium on Operating Systems Design and Implementation (OSDI);[34],Scaling distributed machine learning with the parameter server,2014,USENIX Symposium on Operating Systems Design and Implementation (OSDI);[59],Horovod: fast and easy distributed deep learning in tensorflow,2018,无URL;[66],Optimization of collective communication operations in MPICH,2005,The International Journal of High Performance Computing Applications】。 本节将探索稀疏张量同步方案的设计空间。给定张量,其同步结果是具有相同索引的梯度被聚合,并且所有GPU具有相同的聚合结果。我们将讨论构建稀疏张量同步方案的四个维度。

Figure 5: 具有Hierarchy的两种聚合模式的说明。每个GPU上的梯度来自同一参数,4.7是最终聚合结果。
Figure 5: 具有Hierarchy的两种聚合模式的说明。每个GPU上的梯度来自同一参数,4.7是最终聚合结果。

通信维度。 同步通常有三种通信模式:1) Ring,2) Hierarchy,和3) Point-to-point。它们在图4中以四个GPU为例说明,其中GPU P3聚合所有GPU的数据。在Ring中,所有GPU形成环结构。P0首先将其数据发送到P1,然后P1将数据连同自己的数据传递到P2,依此类推,直到P3接收所有数据。在Hierarchy中,所有GPU形成分层结构,P3是根。在图4b中有两个阶段。在第一阶段,P0将其数据发送到P1,P2将其数据发送到P3。在第二阶段,P1将自己和P0的数据发送到P3。在Point-to-point通信中,其他三个GPU直接将数据发送到P3。

聚合维度。 通信模式可能有多个通信阶段,有两种聚合选项:1) 增量聚合,即在每个阶段聚合张量;和2) 一次性聚合,即仅在最后阶段后聚合所有GPU的张量。这两种选项由于§2.1中讨论的C1和C2导致不同通信阶段的流量不同。在图4中,Ring有三个阶段,Hierarchy有两个阶段。图5显示了以Hierarchy为通信模式的示例。当P1从P0接收张量时,它有两个张量。P1可以聚合两个张量并将聚合结果发送到P3,如图5a所示;或者它可以仅将连接的张量发送到P3,如图5b所示。

Figure 6: Point-to-point的两种分区模式的说明。在(a)中,每个张量作为整体通信,每个GPU接收所有张量。在(b)中,每个张量分割成三个分区;不同GPU的相同分区发送到同一地方,然后聚合结果发送回所有GPU。
Figure 6: Point-to-point的两种分区模式的说明。在(a)中,每个张量作为整体通信,每个GPU接收所有张量。在(b)中,每个张量分割成三个分区;不同GPU的相同分区发送到同一地方,然后聚合结果发送回所有GPU。

分区维度。 有两种分区模式来确保同步后所有GPU具有相同的聚合结果:1) 集中化,其中每个张量作为整体通信和聚合;和2) 并行化,其中每个张量分解成多个分区,每个分区单独通信和聚合。图6比较了以Point-to-point为通信模式的两种分区模式。在集中化中,如图6a所示,每个GPU将其张量作为整体发送到其他GPU。在并行化中,如图6b所示,每个GPU首先将其张量分解成三个分区,需要两个步骤进行同步。第一步在不同地方聚合不同GPU的相同分区,第二步确保所有GPU具有所有分区的聚合结果。

Figure 7: Point-to-point和Parallelism的平衡模式的说明。每个GPU有六个非零梯度,数字是它们的索引。在(a)中,每个GPU的四个梯度发送到GPU 1。在(b)中,每个GPU发送两个梯度到其他GPU,通信平衡良好。然而,实现这种平衡通信是非平凡的。
Figure 7: Point-to-point和Parallelism的平衡模式的说明。每个GPU有六个非零梯度,数字是它们的索引。在(a)中,每个GPU的四个梯度发送到GPU 1。在(b)中,每个GPU发送两个梯度到其他GPU,通信平衡良好。然而,实现这种平衡通信是非平凡的。

平衡维度。 在并行化中,每个分区的非零梯度数量可能变化。因此,在每个GPU接收的流量方面有两种模式:1) 平衡通信,其中GPU接收相同量的数据;和2) 不平衡通信,其中不同GPU接收的流量大大不同。图7比较了三个GPU的Point-to-point下的两种平衡模式。张量中有15个梯度,其中六个是非零的。如图7a所示,中间分区中有四个非零梯度,它们发送到GPU 1。GPU 1接收的流量是GPU 0和GPU 2的4倍。在图7b中,每个GPU发送两个非零梯度到其他GPU,体积平衡良好。由于§2.1中讨论的C3,天真分区会导致不平衡通信。

四个维度描述了稀疏张量同步方案的设计空间。 表3根据它们的维度分类现有方案【[23],Efficient sparse collective communication and its application to accelerate distributed deep learning,2021,ACM SIGCOMM 2021 Conference;[35],PyTorch Distributed: Experiences on accelerating data parallel training,2020,Proceedings of the VLDB Endowment;[56],SparCML: High-performance sparse communication for machine learning,2019,International Conference for High Performance Computing, Networking, Storage and Analysis】。

2.3 最优同步方案

接下来,我们基于四个设计维度分析稀疏张量的最优同步方案。 为方便,我们引入两个特殊方案:平衡并行化和分层集中化。

定义6(平衡并行化)。 它指由[Point-to-point、增量聚合、并行化和平衡通信]表征的同步方案。

定义7(分层集中化)。 它指由[Hierarchy、增量聚合和集中化]表征的同步方案。

定理1(最优方案)。 为了最小化稀疏张量的通信时间,最优同步方案是平衡并行化或分层集中化。

证明。 定理1使用引理1和引理2证明。这里,我们呈现主要直觉,详细证明在附录B.1。

引理1。 当分区模式固定为并行化时,最优方案是平衡并行化。

引理1有三个直觉: 1) 平衡通信优于不平衡通信;2) Point-to-point通信通过最小化并行化分区模式中唯一梯度的流量优于Ring和Hierarchy;3) 增量聚合通过聚合稀疏张量的重叠减少流量优于一次性聚合。

引理2。 当分区模式固定为集中化时,最优方案是分层集中化。

引理2有两个直觉: 1) 当分区模式固定为集中化时,搜索空间减少到六个候选,因为平衡维度不适用。2) 在这些候选中,分层集中化最小化重叠梯度的流量。我们用极端情况演示第二个直觉。假设具有索引idx的非零梯度出现在所有稀疏张量中。在Point-to-point或一次性聚合中,每个GPU必须接收此梯度n-1次。在Ring中,每个GPU的梯度在每个阶段聚合并转发到下一个GPU,导致每个GPU接收梯度n-1次。然而,使用Hierarchy和增量聚合,每个GPU仅接收梯度$\log n$次。

引理1和引理2蕴含定理1。

通信时间分析。 我们将分析定理1中两个方案的理论通信时间。分析中使用的符号列在表4。我们假设每个节点配备一个GPU,并且每对节点有直接双向连接【[56],SparCML: High-performance sparse communication for machine learning,2019,International Conference for High Performance Computing, Networking, Storage and Analysis】。理论通信时间定义为消息传输时间$L/b$,其中L是消息大小,b是网络带宽。为此分析,我们采用COO稀疏格式。为简单起见,我们假设每个GPU上的张量具有相同的$d_G$,所有GPU的平均致密化比率为$\gamma_n$。此外,GPU数量n是2的幂。

平衡并行化。 此方案中的通信涉及两个步骤。第一步中每个GPU接收的流量是$2(n-1)/n \cdot d_G M$,第二步是$2(n-1)/n \cdot \gamma_n^G d_G M$。每个GPU接收的总流量是$2(n-1) (\gamma_n^G + 1) d_G M / n$,导致通信时间为:

$T_{BP} = \frac{2(n-1)(\gamma_n^G + 1) d_G M}{n b}$

分层集中化。 此方案在$\log n$个阶段进行。在第i阶段,每个GPU接收的流量是$2 \gamma_{2^{i-1}}^G d_G M$。每个GPU接收的总流量是$2 \sum_{i=1}^{\log n} \gamma_{2^{i-1}}^G d_G M$。相应的通信时间是:

$T_{HC} = \frac{2 \sum_{i=1}^{\log n} \gamma_{2^{i-1}}^G d_G M}{b}$

平衡并行化作为实际最优方案。 根据定理1,最优方案通过比较方程(1)和(2)确定。基于DL模型中稀疏张量的特性,我们观察到平衡并行化通常优于其他方案,因为在实际分布式训练中$\sum_{i=1}^{\log n} \gamma_{2^{i-1}}^G > (n-1)/n (\gamma_n^G + 1)$。为说明,我们考虑两个极端情况。第一种情况是任意两个张量完全重叠,即$\gamma_n = d_G$。左侧项变为$\log n$,右侧项是$2(n-1)/n < 2$,后者小得多。对于典型的大n,如n≥16,右侧项小几倍。第二种情况是任意两个张量无重叠,即$\gamma_k = k d_G$。左侧项变为n-1,右侧项是(n-1)(1 + 1/n),略大于n-1。然而,实际重叠比率远高于0.05,如图2a所示。更具体地说,当n=16时,即使任意两个张量的重叠比率低至0.05,也满足$\sum_{i=1}^{\log n} \gamma_{2^{i-1}}^G > (n-1)/n (\gamma_n^G + 1)$。

表3:基于维度的不同稀疏张量同步方案比较。

表3图片
表3图片

表4:本文使用的符号。

表4图片
表4图片

2.4 NMT模型的案例研究

在本节中,我们将经验验证定理1。 如表3所列,提出了几种稀疏张量同步方案【[23],Efficient sparse collective communication and its application to accelerate distributed deep learning,2021,ACM SIGCOMM 2021 Conference;[35],PyTorch Distributed: Experiences on accelerating data parallel training,2020,Proceedings of the VLDB Endowment;[56],SparCML: High-performance sparse communication for machine learning,2019,International Conference for High Performance Computing, Networking, Storage and Analysis】。我们从算法视角比较它们的性能,即仅考虑理论通信时间,忽略其他开销,如聚合计算时间和稀疏张量编码解码开销。

AGsparse【[35],PyTorch Distributed: Experiences on accelerating data parallel training,2020,Proceedings of the VLDB Endowment】。 它采用[一次性聚合、集中化],并分别收集非零梯度和相应索引。它不能充分利用稀疏张量间的重叠减少流量。注意,不同实现有不同通信模式【[66],Optimization of collective communication operations in MPICH,2005,The International Journal of High Performance Computing Applications】。

SparCML【[56],SparCML: High-performance sparse communication for machine learning,2019,International Conference for High Performance Computing, Networking, Storage and Analysis】。 它采用分层集中化。它不能充分利用稀疏张量间的重叠减少流量。AGsparse和SparCML的性能取决于重叠。重叠越少,性能越接近最优。然而,如图2所示,DL模型中跨GPU的稀疏张量有显著重叠。

OmniReduce【[23],Efficient sparse collective communication and its application to accelerate distributed deep learning,2021,ACM SIGCOMM 2021 Conference】。 它采用[Point-to-point、一次性聚合和并行化]。OmniReduce由工作者和聚合器组成。它将梯度张量分割成梯度块,仅发送非零块,即至少有一个非零梯度的块,到聚合器进行聚合。OmniReduce不需要传输非零梯度的索引。它需要多个聚合器以获得更好可扩展性。然而,由于均匀分区张量,其性能受不平衡通信影响。

稀疏数据格式按每个方案提出,即AGsparse和SparCML使用COO;OmniReduce使用张量块格式。 我们假设平衡并行化的稀疏数据格式为COO以公平比较。

平衡并行化最优的情况。 图8比较了这些同步方案在批次大小为64的NMT中同步稀疏张量的性能。其他DL模型的比较有类似趋势。它们的通信时间归一化到AllReduce,这是稠密张量的同步方案。AGsparse的通信时间随GPU数量线性增加。它在超过40个GPU时劣于AllReduce,因为它不利用稀疏张量间的重叠减少流量。OmniReduce在少量GPU时优于AllReduce。然而,其性能改进在超过64个GPU时变得边缘。由于非零梯度的偏斜分布,大多数非零梯度在一个分区中,导致不平衡通信。此外,张量在聚合后变得更稠密。当将此分区分割成张量块(例如,每个块有256个梯度【[23]】)时,大多数是非零块。因此,此分区中几乎所有梯度发送到一个聚合器,成为通信瓶颈。SparCML在大量GPU时劣于AllReduce,因为每个GPU接收重复索引及其梯度。相比之下,平衡并行化大大优于现有同步方案和AllReduce。例如,现有方案在128个GPU时不能比AllReduce减少通信时间,但平衡并行化的通信时间仍比AllReduce低36%。

SparCML最优的情况。 这种情况在实际分布式训练中罕见。我们必须将NMT的批次大小减小到很小值以演示。例如,当批次大小为1且GPU数量为8时,嵌入层的梯度张量密度小于0.1%,稀疏张量几乎无重叠,即$\gamma_k \approx k$。在此场景中,SparCML分别优于AGsparse和平衡并行化4%和9%。随着GPU数量增加,SparCML仍优于平衡并行化,但性能差距变得更边缘。

Figure 8: 不同方案在批次大小为64的NMT [40]中同步稀疏张量的比较。
Figure 8: 不同方案在批次大小为64的NMT [40]中同步稀疏张量的比较。

索引通信开销不可忽略。 当稀疏数据格式为COO时,稀疏张量中的每个非零梯度带有索引,显著增加梯度同步的流量和通信时间。为演示其高昂开销,图8显示了一个理想方案名为平衡并行化无索引,其中仅用平衡并行化同步非零梯度而无任何索引信息。相比理想方案,平衡并行化使通信时间加倍。

A2 方法细节

3 ZEN系统

我们将上述观察和洞见结晶成一个整体梯度同步系统ZEN(图1),它利用DL模型中的稀疏性来最小化分布式训练中的同步时间。 ZEN包括定理1中描述的两个方案,以在不同场景下最小化通信时间。在运行时,ZEN收集前几个迭代的轻量稀疏性剖析结果,并通过比较方程(1)和(2)智能确定最优方案。

我们将在本节聚焦平衡并行化,因为目前没有现有解决方案实现它,而SparCML【[56],SparCML: High-performance sparse communication for machine learning,2019,International Conference for High Performance Computing, Networking, Storage and Analysis】已被提出。 我们首先公式化平衡并行化的问题(§3.1),然后开发分层散列算法来解决它(§3.2)。然后我们设计一种新稀疏数据格式来最小化稀疏张量中梯度索引引起的通信开销(§3.3)。

3.1 平衡并行化公式化

为方便,我们从参数服务器(PS)架构【[34],Scaling distributed machine learning with the parameter server,2014,USENIX Symposium on Operating Systems Design and Implementation (OSDI)】借用工作者和服务器的概念到平衡并行化。 因为平衡并行化中有两个梯度同步步骤,我们也称它们分别为Push和Pull。

假设平衡并行化中有n个工作者和n个服务器。 $I_i \subset \mathbb{N}^+$是工作者i生成的非零梯度索引。我们定义实现平衡并行化的问题如下。

问题1。 令I表示$\{I_1, I_2, \cdots, I_n\}$的并集。我们希望有一个映射$f : I \to [n]$使得:

  1. 对于每个$i \in [n]$和$j \in [n]$,集合$\{idx \in I_i | f(idx) = j\}$的基数等于$|I_i|/n$。

  2. 对于每个$j \in [n]$,集合$\{idx \in I | f(idx) = j\}$的基数等于$|I|/n$。

这里我们相应阐述映射f的两个要求如下。

  1. Push中的负载平衡。 对于每个工作者,映射f需要将其非零梯度均匀分解成n个分区。因此,工作者可以将相同量的非零梯度传输到每个服务器。

  2. Pull中的负载平衡。 每个服务器在聚合后应具有相同数量的非零梯度。它还意味着不同工作者的相同索引应发送到同一服务器。

定义8(不平衡比率)。 给定映射f将$I_i$分解成n个分区,表示为$\{I_i^1, \cdots, I_i^n\}$。f的Push不平衡比率是$\max_{i,j \in [n]} \{ n |I_i^j| / |I_i| \}$。

令I表示$\{I_1, I_2, \cdots, I_n\}$的并集,聚合后n个服务器上的索引集为$\{I_1', I_2', \cdots, I_n'\}$。 f的Pull不平衡比率是$\max_{i \in [n]} \{ n |I_i'| / |I| \}$。

基于定义8,平衡并行化中Push和Pull的不平衡比率是1。 我们的目标是最小化任何非零梯度分布的不平衡比率。

数据相关解决方案导致高昂开销。 由于不同工作者有不同索引集,数据相关解决方案需要分析它们的整体分布并为所有工作者计算一个映射,不可避免地导致不可忽略的计算开销。在我们的测试床中,测量的计算成本比迭代时间大几个数量级。因此,我们无法为每个迭代应用数据相关算法并获得映射f。一个可能的方法是定期计算f并维持它到下个迭代。然而,由于迭代间索引分布变化,此方法仍导致高不平衡比率。一个遵循此方法的稻草人是排序索引集I,均匀分区成n部分,并使用边界索引作为阈值来分区下个迭代的索引集。当为NMT模型【[40],Effective approaches to attention-based neural machine translation,2015,Conference on Empirical Methods in Natural Language Processing】每1000迭代计算阈值(n=16)并应用于后续迭代时,Push的不平衡比率在1.4和5.1之间波动,导致服务器间不平衡通信。而且,数据相关解决方案引入的不平衡通信使迭代时间难以估计。许多GPU集群的资源调度机制假设可预测稳定的迭代时间来分配资源给训练作业【[41],Themis: Fair and efficient GPU cluster scheduling,2020,17th USENIX Symposium on Networked Systems Design and Implementation (NSDI 20);[48],Heterogeneity-aware cluster scheduling policies for deep learning workloads,2020,14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20);[54],Optimus: an efficient dynamic resource scheduler for deep learning clusters,2018,Thirteenth EuroSys Conference】。使用波动通信时间调度GPU资源很麻烦。

3.2 数据无关的分层散列

由于数据相关解决方案的局限性,我们必须开发数据无关解决方案,以在Push和Pull中实现负载平衡,同时计算开销可忽略。

解决问题1的一个天真解决方案是在GPU上的多个线程应用通用散列函数,即每个线程独立对不相交输入操作散列函数并写入散列内存。 不幸的是,此方法有损失。当两个索引散列到散列内存的同一位置时,仅一个索引可写入,而另一个被覆盖,导致显著信息损失。例如,当散列内存大小等于密度为10%的稠密张量大小时,由于散列冲突,17.5%的梯度丢失。增加散列内存大小可减少信息损失,但引入不可忽略的计算开销,因为算法必须在散列操作后从散列内存提取非零值。为说明,考虑LSTM【[44],Regularizing and optimizing LSTM language models,2017,无URL】中的嵌入层(1.6GB)。当散列内存大小是张量大小的四倍时,信息损失率降至4.8%。然而,使用PyTorch 2.2中的内置nonzero() API在NVIDIA V100 GPU上的提取开销增加到41.8ms。此开销不可接受,因为它大约是单GPU迭代时间(我们的测试床中114ms)的40%。此外,增加散列内存大小显著增加GPU内存使用,可能导致内存不足并崩溃训练过程。

我们开发了一种新型分层算法来解决问题1。 我们利用GPU上的多个线程高效执行散列函数。我们通过四种技术实现无信息损失和最小GPU内存使用。

技术#1:面向通信的散列内存管理。 散列内存分成多个分区,每个分区的数据写入对应服务器。每个分区进一步分成并行内存和串行内存以防止数据丢失。当GPU内的线程执行散列函数时,它首先检查散列冲突,即检查散列位置是否被占用。如果无冲突,数据写入并行内存,启用所有线程的并发写入。如果有冲突,冲突索引使用原子操作顺序写入串行内存,仅允许一个线程一次写入。尽管内存中的索引顺序随机,但无需排序,因为它们的顺序不影响聚合结果。

不幸的是,我们观察到当散列冲突率升高时,串行写入的开销变得显著。

技术#2:GPU中每个线程的多个散列函数。 使用多个散列函数减少串行写入成本。当发生散列冲突时,GPU线程可以用新散列函数重新散列索引到并行内存中的另一个位置。有机会这个新位置可用,减少串行内存的顺序写入操作数。尽管多个散列函数仍有散列冲突,我们观察到使用四个散列函数的冲突率小于1%,串行写入串行内存的成本变得可接受。

然而,使用多个散列函数的重新散列可能导致不完整聚合。 因为不同GPU有不同非零梯度索引集,它们的散列索引序列也不同。因此,特定索引的位置在GPU间可能不同。例如,两个索引idx1和idx2,其中idx1 < idx2,用第一个散列函数散列到同一位置。GPU 1有idx2;GPU 2有idx1和idx2。在GPU 1中,idx2的位置由第一个散列函数确定,但在GPU 2中,idx2的位置由第二个散列函数确定,因为第一个散列的位置已被idx1占用。随后,分区内存将导致相同索引在不同GPU分配到不同分区,导致不完整聚合。

技术#3:工作者间的consistent分层散列。 我们提出两级分层散列算法来保证完整聚合。第一级散列确定索引所属的分区,并保证索引在所有GPU中属于同一分区。第二级散列确定其在此分区中的位置。为确保不同工作者的相同索引可发送到同一服务器,ZEN为所有工作者分配相同的第一级散列函数,但允许它们有独立的第二级散列函数。

技术#4:无锁读后写机制。 有一个担忧是两个线程的两个值可能同时散列到同一可用位置,这种散列冲突不能被常规机制检测,导致信息损失。ZEN使用读后写机制检查此冲突,消除此类信息损失。在内存写入后,线程读取位置存储的值,此操作无依赖于其他线程的值。如果值等于它写入的,线程继续下一个输入。然而,如果值不是它写入的,它意味着有散列冲突,其值被覆盖。然后线程进行重新散列或串行写入。

分层散列算法。 带有四个技术的伪代码显示在算法1中。此算法的示例也如图9所示。给定稠密张量G及其非零梯度索引I,它分配形状为n × (r1 + r2)的内存x,其中n是分区数,r1是并行散列操作的内存大小,r2是串行内存大小。它并行为每个idx ∈ I执行散列操作(行4-17)。通用散列函数h0 : N+ → [n]用于将idx定位到分区p = h0(idx)(行5)。算法还需要k个通用散列函数H = {h1, · · · , hk},其中hi : N+ → [r1]。确定分区p后,算法尝试用h1找到可用目的地x[p][h1(idx)]。如果此位置可用,idx写入其中。否则,算法用h2重新散列idx找到新位置。它最多重新散列索引k轮,使用hi作为第i轮的散列函数,直到找到可用目的地(行6-16)。k次重新散列尝试失败后,算法将idx写入分区p的串行内存(行8-11)。串行写入是原子操作(行9-10)以确保无信息损失。一旦所有索引写入内存,它从内存提取稀疏张量(行19-23)。

接下来,我们将分析算法1的属性及其不平衡比率。

每个迭代对模型准确性无影响。 ZEN是稀疏张量同步的高效通信方案,它对模型的迭代收敛率无影响。对于具有自然张量稀疏性的分布式训练,ZEN的稀疏张量聚合结果与AllReduce的相应稠密张量相同。当使用稀疏化算法时,由于压缩本身可能发生一些准确性损失,取决于稀疏水平。然而,ZEN不引入超出稀疏化算法固有的额外准确性损失。具体地说,ZEN的Push通过其四个技术确保无进一步信息退化,其梯度聚合和Pull操作都是无损的。ZEN的稀疏同步后聚合张量与AGsparse【[35],PyTorch Distributed: Experiences on accelerating data parallel training,2020,Proceedings of the VLDB Endowment】和OmniReduce【[23],Efficient sparse collective communication and its application to accelerate distributed deep learning,2021,ACM SIGCOMM 2021 Conference】等其他方案产生相同。因此,ZEN保留相同的每个迭代模型准确性,同时实现更短的壁钟迭代时间。更详细的准确性评估和讨论如图18所示。

可忽略的内存和提取开销。 得益于多个散列函数和串行内存,算法1可以用小内存大小实现无信息损失,在我们的实现中通常小于150MB(<1% GPU内存容量)用于所有评估模型。从散列后内存提取索引(算法1中的行20)的开销变得可忽略。

无依赖于工作负载。 注意,我们不对数据分布做假设,仅使用通用散列算法的属性。因此,算法1可以为DL训练工作负载的不同非零梯度分布获得一般理论保证。

保证的不平衡比率。 因为散列函数h0确定每个索引的分区,算法1的不平衡比率由以下定理保证。

定理2(算法1的负载平衡)。 给定具有|G|参数的稠密张量G。算法1提供映射f : I → [n]使得

  1. 以至少1 - 1/n的概率,其Push不平衡比率至多1 + Θ(√(n log n / (|G| d_G)))。

  2. 以至少1 - 1/n的概率,其Pull不平衡比率至多1 + Θ(√(n log n / (|G| d_n^G)))。

证明可在附录B.2找到。 因为n log n比|G|小几个数量级,算法1对问题1的精确解进行良好逼近,用于Push和Pull。假设n=128,|G|=10^7,且d_n^G=0.5,则√(n log n / (|G| d_n^G)) < 0.02。其实际不平衡比率总是小于1.1用于研究的六个模型。

算法1:分层散列算法伪代码

Input: G is a dense tensor and I  N+ is the indices of its non-zero gradients. n is the number of partitions. Each partition has a memory size r1 + r2, where r1 and r2 are the memory sizes for parallel and serial operations, respectively. h0 : N+  [n] is a universal hash function. H = {h1, · · · , hk} is a set of universal hash functions where hi : N+  [r]. 
Output: The partitioned sparse tensors. 
1 Function hierarchical_hash(I, G, h0, H): 
2     Allocate memory x  0n×(r1+r2) 
3     Allocate atomic counters c  r1n 
4     foreach idx  I in parallel do 
5         p  h0(idx) 
6         for i  1 to k + 1 do 
7             q  hi(idx) 
8             if i = k + 1 then 
9                 q  atomicAdd(c[p],1) 
10                x[ p][q]  idx 
11            end 
12            if x[p][q] == 0 then 
13                x[ p][q]  idx 
14                break 
15            end 
16        end 
17    end 
18    out put = [] 
19    for i  0 to n  1 do 
20        indices = nonzero(x[i]) 
21        values = G[indices] 
22        out put.append((indices, values)) 
23    end 
24    return out put

Figure 9: 分层散列算法的演示。我们对索引进行并行散列。对于每个索引,我们用散列函数h0分配其分区。我们接下来用散列函数h1分配到第一个位置。然而,因为此位置被占用,我们用函数h2重新散列到第四个位置。因为它也被占用,我们用原子操作串行地将索引写入串行内存。
Figure 9: 分层散列算法的演示。我们对索引进行并行散列。对于每个索引,我们用散列函数h0分配其分区。我们接下来用散列函数h1分配到第一个位置。然而,因为此位置被占用,我们用函数h2重新散列到第四个位置。因为它也被占用,我们用原子操作串行地将索引写入串行内存。

3.3 成本高效的编码方案

3.3.1 现有稀疏格式低效

有几种稀疏格式来表示稀疏张量用于同步。 不幸的是,它们都不能最小化非零梯度索引引起的开销。我们假设梯度数据类型是FP32。

COO。 它在低张量密度时高效【[39],Deep gradient compression: Reducing the communication bandwidth for distributed training,2017,The International Conference on Learning Representations (ICLR);[74],DRAGONN: Distributed randomized approximate gradients of neural networks,2022,International Conference on Machine Learning】。然而,它使流量加倍,并在高密度时低效。如图2b所示,张量在聚合后变得更稠密。例如,BERT的平均张量密度从1.06%增加到128个GPU稀疏张量聚合后的40.8%。理论上,在Pull中传输稀疏张量比传输稠密张量减少2.5倍流量,但由于非零梯度索引,减少仅缩减到1.2倍。

张量块。 它在OmniReduce中使用【[23],Efficient sparse collective communication and its application to accelerate distributed deep learning,2021,ACM SIGCOMM 2021 Conference】。稠密张量分割成梯度块,仅传输非零块。然而,它在高张量密度时低效。当将高密度张量分割成张量块(例如,每个块有256个梯度)时,大多数块至少有一个非零梯度,成为非零张量块。同步方案必须通信大量零梯度。我们评估了与张量密度相比使用张量块格式通信的梯度百分比。对OPT-2.7B应用DGC选择前5%梯度。当块大小为256时,通信超过30%梯度,远高于5%的张量密度。

位图。 它仅需一位指示梯度是否为零。不幸的是,常规位图仍引起不可忽略的流量来识别非零梯度。当均匀分区稠密张量G时,每个服务器中非零梯度的索引在[1, |G|]的子范围内。例如,当|G|=15且有三个服务器时,服务器i的索引范围是[5i + 1, 5(i + 1)]。表示每个服务器中非零梯度索引的额外位图大小是|G|/n/32,当梯度数据类型是FP32时。每个工作者接收的总位图大小是|G|/32。当G的张量大小是856MB,等于DeepFM中的嵌入表大小时,总位图大小是27MB。虽然算法1启用负载平衡,但每个服务器中的非零梯度随机分布在整个范围内。如果我们仍用位图表示索引,每个服务器的额外位图大小是|G|/32,每个工作者接收的总位图大小变为n |G|/32,随服务器数量线性增加。当有16个服务器时,总位图大小是428MB。

3.3.2 散列位图

我们为ZEN开发了一种新型散列位图来最小化Pull中表示非零梯度索引的开销。 给定稠密张量G和算法1中的h0,每个工作者中应推送到服务器i的索引集$I_i = \{idx \in [1, |G|] | h_0(idx) = i\}$被确定,虽然不在连续范围内。由于$I_i$和$I_j$在i ≠ j时不相交,它提供基于$I_i$而不是整个范围构建位图的机会。

Figure 10: 散列位图的说明。
Figure 10: 散列位图的说明。

图10说明了散列位图如何为|G|=15和三个服务器的I0工作。 非零梯度的两个索引是{5,7}。hash_bitmap_encode()用于编码索引。给定稠密张量G,它首先根据I0中的索引提取本地梯度。然后将本地梯度编码到位图。因为第二个和第三个梯度是非零的,散列位图中的第二位和第三位是1,其他位是0。hash_bitmap_decode()用于将散列位图解码到索引集。它首先将散列位图解码到位图索引,即1的索引。例如,因为散列位图中的第二位和第三位是1,位图索引是{2, 3}。然后它用位图索引作为索引从I0中提取相应值作为非零梯度的全局索引。在此示例中,值是{5, 7},正好是非零梯度的两个索引。伪代码显示在算法2中。

函数hash_bitmap_encode()在每个服务器上调用,然后将散列位图广播到所有工作者。 每个工作者接收所有服务器的散列位图后,它调用hash_bitmap_decode()用相应$I_i$将散列位图解码到索引。注意,$I_i$在训练开始时离线计算和排序,对于服务器和工作者的相同h0保持不变。

散列位图保证每个工作者接收的总散列位图大小在ZEN的Pull中恒定为|G|/32。 假设有n个服务器。应推送到服务器i的索引集是$I_i = \{idx \in [1, |G|] | h_0(idx) = i\}$。用hash_bitmap_encode(),服务器i中编码的散列位图大小是|I_i|/32。因为每个工作者需要接收所有服务器的散列位图,总大小是$\sum_{i=0}^{n-1} |I_i|/32 = |G|/32$。ZEN在Push中仍使用COO表示稀疏张量,由于低张量密度。

算法2:散列位图伪代码

算法2图片
算法2图片

4 实现

我们用约900行Python代码、250行CUDA代码和500行用于修改ColossalAI【[36],Colossal-AI: A unified deep learning system for large-scale parallel training,2023,52nd International Conference on Parallel Processing】实现ZEN 1。 对于DL模型训练,我们使用ColossalAI实现数据并行和混合并行(数据并行 + 张量并行,张量并行度设为8,通常在单个节点内)。我们还扩展ColossalAI以自定义策略支持Google的Gemma2【[65],Gemma 2: Improving open language models at a practical size,2024,无URL】模型,用于高效张量并行。

对于自然张量稀疏性,我们仅对嵌入层的梯度应用稀疏张量同步方案,用于跨机器通信。 对于稀疏化算法,我们使用DGC【[39],Deep gradient compression: Reducing the communication bandwidth for distributed training,2017,The International Conference on Learning Representations (ICLR)】并适应张量并行:每个设备本地采样数据,跨所有设备收集这些样本,基于聚合数据计算全局top-k阈值,并在每个设备本地应用此阈值确定top-k值。我们在PyTorch的DDP【[35],PyTorch Distributed: Experiences on accelerating data parallel training,2020,Proceedings of the VLDB Endowment】中注册custom_comm_hook以启用DGC和我们的同步方案。梯度计算在反向传播期间与梯度同步重叠。在应用DGC到张量前,梯度张量在DDP中融合成通信桶【[73],Cupcake: A compression scheduler for scalable communication-efficient distributed training,2023,Proceedings of Machine Learning and Systems】。我们确定128MB桶大小最适合评估模型。

我们在CUDA C中实现分层散列算法,并作为PyTorch扩展使用。 使用的散列函数是MurmurHash【[9],Murmurhash 2.0,2008,无URL】。我们为MurmurHash设置不同种子生成不同散列函数。在训练开始时,ZEN生成一组随机种子并广播到所有GPU以确保工作者间的散列一致性。ZEN使用ReduceScatter/AllGather【[25],MPICH2: A new start for MPI implementations,2002,European Parallel Virtual Machine/Message Passing Interface Users’ Group Meeting;[66],Optimization of collective communication operations in MPICH,2005,The International Journal of High Performance Computing Applications】通信机器内的稠密张量,因为GPU机器通常配备NVLink。在我们的评估中,平衡并行化是研究模型的最优方案。

A4 实验环境

数据集名称、规模及用途。 对于自然稀疏性:One Billion Word【[1],Billion Word Benchmark,https://code.google.com/archive/p/1-billion-word-language-modeling-benchmark/】用于LSTM,规模1亿词,用于语言建模;Criteo【[2],Criteo dataset,https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/】用于DeepFM,规模1TB点击日志,用于CTR预测;IWSLT 2014 De-En【[3],IWSLT 2014 De-En dataset,https://sites.google.com/site/iwsltevaluation2014/data-provided】用于NMT,规模标准,用于机器翻译。批次大小分别为128、1024、64。对于梯度压缩:RedPajama【[75],RedPajama: an open dataset for training large language models,2024,无URL】用于LLM,规模标准训练语料,TP组内批次大小1,序列长度1024 tokens。

模型架构关键参数。 自然稀疏性模型见表1:LSTM嵌入维度1000,隐藏单元1000;DeepFM嵌入维度16;NMT嵌入维度512,隐藏单元512。LLM见表2:Llama3.2-3B参数3B,层数28,隐藏维度2304,注意力头12;OPT-2.7B参数2.7B,层数32,隐藏维度2560,注意力头32;Gemma2-2B参数2B,层数26,隐藏维度2048,注意力头8。

硬件配置。 16个p3.16xlarge实例,每个8个NVIDIA V100 GPU(16GB内存),NVLink用于节点内,25Gbps网络连接;另16个p3dn.24xlarge,每个8个V100 GPU(32GB内存),NVLink,100Gbps RDMA网络与EFA。

软件配置。 Ubuntu 20.04,CUDA-11.0,PyTorch-2.2,NCCL-2.7.8,CuPy-11.0,ColossalAI-v0.4.6;p3dn.24xlarge额外aws-ofi-nccl-v1.9.2和libfabric-v1.11.1支持EFA。使用DGC选择前5%梯度。

A4 实验结果

训练吞吐量实验(自然稀疏性模型,25Gbps网络)。 如图11所示,ZEN在处理样本/秒上优于所有基线。在16机器上训练LSTM时,ZEN比SparCML加速1.67倍,比OmniReduce加速2.48倍,比AllReduce加速3.1倍。在DeepFM和NMT中,ZEN比OmniReduce加速1.44倍和1.51倍。随着机器增加,ZEN的优势放大,显示良好可扩展性。

Figure 11: 25Gbps网络中自然张量稀疏性模型的训练吞吐量。
Figure 11: 25Gbps网络中自然张量稀疏性模型的训练吞吐量。

训练吞吐量实验(自然稀疏性模型,100Gbps RDMA网络)。 如图12所示,ZEN在LSTM中比SparCML加速1.44倍,在DeepFM和NMT中比AllReduce加速1.33倍和1.38倍,比OmniReduce加速1.25倍和1.32倍。

Figure 12: 100Gbps RDMA网络中自然张量稀疏性模型的训练吞吐量。
Figure 12: 100Gbps RDMA网络中自然张量稀疏性模型的训练吞吐量。

训练吞吐量实验(梯度压缩模型,25Gbps网络)。 如图13所示,对于Llama3.2-3B,ZEN比OmniReduce加速1.68倍,比SparCML加速2.19倍,比AllReduce加速2.02倍。对于OPT-2.7B和Gemma2-2B,ZEN比AllReduce加速2.10倍和2.04倍,比OmniReduce加速1.66倍和1.61倍。

Figure 13: 25Gbps网络中稀疏化算法张量稀疏性模型的训练吞吐量。
Figure 13: 25Gbps网络中稀疏化算法张量稀疏性模型的训练吞吐量。

训练吞吐量实验(梯度压缩模型,100Gbps RDMA网络)。 如图14所示,ZEN比AllReduce在Llama3.2-3B、OPT-2.7B和Gemma2-2B中加速64%、45%和44%。比OmniReduce加速1.32倍、1.31倍和1.27倍。

Figure 14: 100Gbps RDMA网络中稀疏化算法张量稀疏性模型的训练吞吐量。
Figure 14: 100Gbps RDMA网络中稀疏化算法张量稀疏性模型的训练吞吐量。

通信改进实验(16机器,25Gbps网络)。 如图15所示,OmniReduce加速高达1.58倍。AGsparse、Sparse PS和SparCML有时劣于AllReduce,因为高密度下COO稀疏张量大小大于稠密张量。ZEN在LSTM中加速6.77倍,在Gemma2-2B中加速3.51倍,比SparCML和OmniReduce加速高达2.82倍和5.16倍;在DeepFM和OPT-2.7B中比AllReduce加速2.10倍和3.44倍。

Figure 15: 25Gbps网络中比AllReduce的通信加速。
Figure 15: 25Gbps网络中比AllReduce的通信加速。

算法1参数研究实验(模拟张量214M参数)。 如图16a所示,增加r1从|G| d_G到2|G| d_G显著减少操作成本,但进一步到4|G| d_G增加时间。如图16b所示,k从1到3减少成本,但k=3和4类似。

Figure 16: 算法1的计算开销。(a) 不同内存大小和(b) 不同重新散列数。
Figure 16: 算法1的计算开销。(a) 不同内存大小和(b) 不同重新散列数。

可忽略计算开销实验。 对于DeepFM的214M参数张量,散列开销约6ms,而比AllReduce的通信节省270ms(25Gbps),或占100Gbps节省的9%。

散列位图有效性实验(16服务器)。 如图17a所示,散列位图与COO的差距随密度增加增大,显著优于位图和张量块。在95%密度下仍减少流量,而位图和COO在>50%时不能节省。

Figure 17: (a) 散列位图的有效性;(b) ZEN在25Gbps网络中16机器的性能分解。
Figure 17: (a) 散列位图的有效性;(b) ZEN在25Gbps网络中16机器的性能分解。

性能分解实验(16机器,比AllReduce)。 如图17b所示,主要益处来自算法1,散列位图额外益处明显。例如,用COO后加速2.74倍(LSTM)和1.53倍(Llama3.2-3B),散列位图进一步改进13%和34%。

模型准确性验证实验。 如图18a所示,DeepFM的每个迭代测试准确性与AllReduce完全匹配。如图18b所示,OPT-2.7B的训练损失与AGsparse相同,证明ZEN保留收敛率。

Figure 18: ZEN对每个迭代模型准确性无影响。
Figure 18: ZEN对每个迭代模型准确性无影响。

A5 结论

论文的结论是:系统探索稀疏张量同步方案的设计空间,识别最优方案;提出ZEN,使用并行GPU计算的无损失新型散列算法实现最优同步。ZEN显著提高稀疏DNN模型训练效率,而不对模型准确性有任何影响。未来工作可探索更多硬件加速和稀疏格式优化。

A6 附录

A 稻草人解决方案的算法

在本节中,我们呈现算法3,这是带有直接散列算法的数据无关解决方案的算法。 此算法补充§3.2中的讨论。

算法3:带有散列的稻草人解决方案伪代码

算法3图片
算法3图片

解决问题1的一个天真解决方案是在GPU上的多个线程应用通用散列函数,即每个线程独立对不相交输入操作散列函数并写入散列内存。 给定稠密张量G,其非零梯度索引集是I。算法首先分配形状为n × r的内存x,其中n表示分区数,r是每个分区的内存大小。对于每个idx ∈ I,它使用给定通用散列函数[15] h : N+ → [nr]生成散列值h(idx),其中nr是散列函数h的范围。接下来,它将idx写入分区⌊h(idx)/r⌋中的(h(idx) mod r)位置。散列操作并行执行以最小化计算开销[74]。之后,它从每个分区的内存提取非零索引,并用它们从G中查找相应梯度。最后,它为每个分区返回稀疏张量并推送到相应服务器。

B 理论分析

B.1 定理1的证明

定理1(最优方案)。 为了最小化稀疏张量的通信时间,最优同步方案是平衡并行化或分层集中化。

我们用两个引理证明定理1。

引理1。 当分区模式固定为并行化时,最优方案是平衡并行化。

证明。 有三种通信模式:Ring、Hierarchy和Point-to-point。我们首先考虑带有[Point-to-point通信和并行化]的同步方案,即PS架构。给定n个服务器和密度为d_G的梯度张量G。我们首先分别分析Push和Pull操作的通信时间。然后讨论不同PS方案的通信时间。

Push。 因为偏斜比率是s_n^G,n个分区中最大密度是s_n d_G。从此分区提取的稀疏张量大小是2 s_n d_G M / n。因此,稀疏PS中Push的通信时间是2 (n - 1) s_n d_G M / n / B。

Pull。 聚合后,n个分区中最大密度变为s_n^G d_n^G。在PS架构的现有实现中,Pull的通信时间是2 (n - 1) s_n d_n M / n / B,因为每个服务器需要用Point-to-point通信广播其聚合结果到所有工作者[29, 34]。理论上,有其他方式实现PS架构中的Pull。例如,每个服务器可执行广播集体操作。不同算法的广播性能在[10]中分析,其Pull通信时间可表示为2 b d_n^G M / B,其中b是一个算法中的轮数。例如,当使用Binomial Tree Algorithm时b = ⌈log n⌉,当使用Scatter-AllGather Algorithm时b = 2 (n-1)/n [10, 25]。

稀疏PS。 结合Push和Pull的通信时间与Point-to-point,其整体通信时间是2 (n - 1) (d_G + d_n) s_n M / n / B。

带广播的稀疏PS。 当考虑Pull的广播时,整体通信时间变为2 (n - 1) s_n^G d_G M / n / B + 2 b d_n^G M / B。我们将此情况表示为带广播的稀疏PS。

平衡并行化。 在平衡并行化中,偏斜比率s_n总是1。我们将稀疏PS通信时间中的s_n替换为1,得到平衡并行化的通信时间:2 (n - 1) (d_G + d_n^G) M / n / B。

平衡并行化在并行化方案中最优。 显然,当偏斜比率大时平衡并行化远优于稀疏PS。带广播的Spase PS与平衡并行化的性能比率是s_n^G n b γ_n^G / (s_n^G + b γ_n^G n) / (1 + γ_n^G) (n-1) / (1 + γ_n^G) 1。因为s_n和b都大于1,比率也大于1。因此,平衡并行化总是优于稀疏PS和带广播的稀疏PS在通信时间上。

然后我们证明平衡并行化优于其他方案。 因为一次性聚合不能利用稀疏张量间的重叠,带有一次性聚合的同步方案性能劣于带有增量聚合的。因此,我们仅考虑增量聚合用于Ring或Hierarchy通信的方案。我们考虑它们的最佳情况,即分区后偏斜比率是1。此外,我们仅需比较第一步,因为它们在第二步有相同通信时间。平衡并行化中第一步的通信时间是2 (n - 1) d_G M / n / B。

带有Ring和增量聚合的方案。 它们有n-1个通信阶段。第i阶段的张量密度是d_i^G。注意d_1^G = d_G。因此,通信时间是2 ∑{i=1}^{n-1} d_i^G M / n / B。因为张量在聚合后可能更稠密,当i < j时d_i^G ≤ d_j^G,且∑ d_i^G ≥ (n - 1) d_G。因此,带有Ring和增量聚合方案的通信时间不小于平衡并行化的。}^{n-1

带有Hierarchy和增量聚合的方案。 它们有log n个通信阶段。因为每个分区有分层结构,第i阶段的流量是d_{2^{i-1}}^G 2^{i-1} M / n,总流量在所有log n阶段是V = ∑{i=1}^{log n} d M / n = 2 (n - 1) d_G M。因此,每个GPU接收的流量不小于2 (n - 1) d_G M / n,它们不优于平衡并行化。□}}^G 2^{i-1} M / n。因为d_{2^{i-1}}^G ≥ d_G,V ≥ ∑_{i=1}^{log n} d_G 2^{i-1

引理2。 当分区模式固定为集中化时,最优方案是分层集中化。

证明。 当任意两个稀疏张量无重叠时,每个GPU必须接收的其他GPU的所有张量是最小流量。任何带有集中化的同步方案实现最优通信时间。

当稀疏张量重叠时,令n表示GPU数量,I_0, I_1, ..., I_{n-1}分别是每个GPU中非零梯度索引集。 C是所有稀疏张量的重叠,即C = ∩ I_i。如果同步方案采用Point-to-point通信或一次性聚合,每个GPU必须接收C n-1次。然后我们考虑增量聚合且通信模式是Ring或Hierarchy。在Ring中,每个GPU的张量在每个阶段聚合然后转发到下一个GPU。因此,此张量被每个GPU接收。因为C是共同重叠,每个GPU也必须接收C n-1次。当通信模式是Hierarchy时,每个GPU用其自己的具有log n + 1级的分层结构从所有其他GPU接收数据,如图4b所示。因为根GPU在每个级中,它必须接收C log n次。它表明带有[Hierarchy、增量聚合和集中化]的方案的流量小于其他带有集中化的方案。令C'表示稀疏张量子集的重叠,我们可得出类似结论,即GPU子集必须多次接收C'。换句话说,每个GPU仍必须多次接收重叠。□

引理1和引理2蕴含定理1。

B.2 定理2的证明

定理2(算法1的负载平衡)。 给定具有|G|参数的稠密张量G。算法1提供映射f : I → [n]使得

  1. 以至少1 - 1/n的概率,其Push不平衡比率至多1 + Θ(√(n log n / (|G| d_G)))。

  2. 以至少1 - 1/n的概率,其Pull不平衡比率至多1 + Θ(√(n log n / (|G| d_n^G)))。

证明。 算法1的不平衡比率仅由h0 : N+ → [n]确定,而散列函数集H聚焦于精确内存写入。

I_i中的索引数是|G| d_G,其中d_G是G的密度。 因为h0是数据无关的,问题1中的部分1可表述为:给定|G| d_G个球,我们希望用通用散列函数h0将它们扔到n个箱中。从[11]的结果,箱的最大负载至多|G| d_G / n + Θ(√(|G| d_G log n / n)) 以至少1 - 1/n的概率。回忆定义8中Push的不平衡比率:

不平衡比率 = max_{i,j ∈ [n]} { n |I_i^j| / |I_i| }

因为max{|I_i^j|} ≤ |G| d_G / n + Θ(√(|G| d_G log n / n)),我们有

不平衡比率 ≤ 1 + Θ(√(n log n / (|G| d_G)))

以至少1 - o(1)的概率。因此,我们完成第一部分的证明。

因为h0是数据无关的,问题1中的部分2可表述为:给定|I| = |G| d_n^G个球,我们希望用通用散列函数h0将它们扔到n个箱中。 箱的最大负载至多|G| d_n^G / n + Θ(√(|G| d_n^G log n / n)) 以至少1 - 1/n的概率。回忆定义8中Pull的不平衡比率:

不平衡比率 = max_{i ∈ [n]} { n |I_i'| / |I| }

因为max{|I_i'|} ≤ |G| d_n^G / n + Θ(√(|G| d_n^G log n / n)),我们有

不平衡比率 ≤ 1 + Θ(√(n log n / (|G| d_n^G)))

以至少1 - 1/n的概率。因此,我们完成第二部分的证明。