应对大规模训练中的通信挑战:同步与拓扑

发布时间:2022-01-24 07:00

摘要

 

随着各个领域数据的不断积累和规范化,越来越多的研究者开始着手解决更大规模的问题,不可避免地会提出或使用更多大规模的模型,分布式训练几乎成为了一个必选项。分布式系统面临的一大挑战是通信问题,特别是对于AI训练任务,使用的是高算力的设备,如果互联是低速的,那么多台设备组成的分布式系统性能甚至会低于单台设备。随着越来越多的大规模训练需求的产生,如何更好地应对分布式训练系统的通信问题已成为研究热点。本文首先介绍分布式训练中的通信问题,然后从减少用于同步的通信和采用更利于通信的拓扑结构两个方面介绍应对大规模训练中通信问题的方法。

 

背景介绍

 

OpenAI在2018年的分析报告[1]中就已提出最大的AI训练任务所需算力每3.4个月翻一番的趋势,从2012年到2018年,人们对于算力需求的增长超过了300,000倍。充足的算力支持可以大大缩短AI研究的进程,比如BERT在16个TPUv3上的预训练时间为81.4个小时,约3天,而在1024个TPUv3上预训练时间可以缩短到76分钟[9]。另一方面,从图1展示的大模型参数量变化趋势可以看出,AI训练任务对存储容量的需求也呈指数增长。此外,随着计算机视觉、自然语言处理、自动驾驶、生物医学等各个领域数据的不断积累和规范化,越来越多的研究者开始着手解决更大规模的问题,不可避免地会提出或使用更多大规模的模型,分布式训练几乎成为了一个必选项。

 

 

图1大模型参数量变化趋势

 

 

分布式系统面临的一大挑战是通信问题,特别是对于AI训练任务,使用的是高算力的设备(比如GPU、TPU),如果互联是低速的(比如1/10 Gbps Ethernet),那么多台设备组成的分布式系统性能甚至会低于单台设备[10]。一个直观的例子如图2,这里假设通信时间是计算时间的25倍,由于通信开销的存在,将多个设备组合成分布式系统时,实际所能带来的速度提升与理想情况下能提供的速度提升之间存在明显的差距。随着越来越多的大规模训练需求的产生,如何更好地应对分布式训练系统的通信问题已成为研究热点。

 

图2 存在通信开销的情况下系统的可扩展性(b: batch size, H: communication interval)[6]

 

 

分布式训练中的通信问题

 

在分析通信问题之前,我们首先梳理一下模型训练的主要流程。模型训练的典型工作流如图3,一次迭代的主要步骤为1)采样小批次数据(mini-batch);2)读入数据,进行前向过程(Feed-Forward),计算目标函数的损失值;3)进行反向过程(Backward Propagation),为模型参数计算梯度;4)根据优化器算法更新(Update)模型参数。在分布式系统中,模型梯度或参数会分布在不同的设备上,故在更新模型参数前,需对不同设备上的梯度或参数进行聚合操作(Gradient/Model Aggregation)。

 

 

图3 模型训练的典型工作流(修改自[2])

 

 

接下来看看分布式训练系统中哪些地方需要通信。常见的并行方式有两种:数据并行和模型并行。

数据并行的一个例子如图4,假设我们有4台设备(Worker 1-4),每台设备均能容纳模型参数,且剩余容量能够支持mini-batch的训练,那么我们就能在每台设备上保存一份相同的模型副本,训练时每台设备负责处理不同的mini-batch,由此会产生不同的梯度,系统会将不同设备产生的梯度聚合到一起,计算均值,再将梯度分发给各个设备进行模型参数更新。故数据并行模式下通信开销主要来源于梯度的聚合和分发。而当单个设备无法支持整个模型的训练时,我们就需要新的并行方案:模型并行。

 

 

图4 数据并行[2]

 

 

模型并行的两个例子如图5,图中黑色双向箭头表示有通信产生,由于在这一模式下,模型同一层的参数会分布在不同的设备上,故前向和反向过程中都需要接收其他设备产生的结果,这就引入了巨大的通信量。此外,每个设备得到的梯度同样需要聚合后再分发给各个设备进行参数更新,同样会产生通信开销。不同的模型划分方法会带来不平衡的参数分布以及不同的计算依赖关系,如何进行合理的模型划分是一个NP完备问题[3]。

 

 

图5 模型并行[2]

 

由于模型训练使用的都是具有强大算力的设备,上述提到的通信需求都是十分高频的需求,当前互联技术所能提供的带宽很难直接满足训练过程中设备间的数据传输。为应对这一通信挑战,我们可以从四个方面入手:1)放宽同步要求,减少用于同步的通信;2)采用更利于通信的拓扑结构;3)对通信数据进行压缩,减少通信数据量;4)通过调度来使通信过程和计算过程尽量并行。限于篇幅,本文先介绍前两个方面的相关技术。

 

减少用于同步的通信

 

在使用多个设备进行分布式训练时,需要大量的同步操作来保证模型的全局一致性。异步设计可以有效提高整个系统的运行效率,但会影响模型的收敛性。下面以数据并行为例来分析四种常见的模型同步模式。

 

 

图6 四种同步模式的工作流对比[2]

 

1. Synchronous

Synchronous模式下,在每次模型更新前,需要完成所有设备的信息收集和聚合,形成全局一致的更新数据,再由各个设备根据统一的更新数据进行更新。这一模式的典型代表是BSP-SGD(Bulk Synchronous Parallel SGD)[4],其工作流如图6中BSP一行,假设有4台设备,每台设备执行前向和反向过程所花费的时间不一,在不同的时间节点产生了梯度,BSP会通过一个Barrier来等待所有设备提供梯度,并完成梯度聚合,再将梯度分发给各个设备进行参数更新,或者由某个设备统一更新参数后分发给各个设备进行同步,各个设备获得最新参数后独立开启下一轮迭代,由此往复。这一设定下每次迭代所花费的时间由最慢的设备决定,且设备数越多,梯度聚合过程通信开销越大,严重影响系统的可扩展性(scalability)。

 

2. Stale-Synchronous

SSP(Stale-Synchronous Parallel)[5]模式允许工作较快的设备提前触发更新操作,其工作流如图6中SSP一行,在完成全局梯度聚合前,运行较快的设备使用本地产生的梯度进行更新,通过一个Staleness bounded Barrier来保证运行较快的设备更新次数与运行较慢的设备不会差距过大。具体做法是每个设备记录迭代次数,与最慢的设备差距到达一定阈值时暂停迭代,等待所有设备产生的梯度完成聚合,进行统一的参数更新。

 

3. Local SGD

Local SGD[6]通过模型聚合来进行同步,其工作流如图6中Local SGD一行,每台设备先自行迭代一定步数,再统一对所有设备上的模型进行平均来获得全局一致的模型。可以看到设备单独迭代的次数越多,通信间隔就越长,用于通信的时间占比就越小。但单独迭代次数越多,不同设备上的模型差异越大,最终收敛性就越差。单独迭代次数的设置是在通信效率和模型收敛性上的一个权衡。

 

4. Asynchronous

ASP(Asynchronous Parallel)[7]模式进一步放宽同步要求,其工作流如图6中ASP一行,每台设备产生梯度后直接进行聚合操作,根据当前聚合的信息进行更新。当设备更多时,这一规则可以推广为一定数量的设备产生梯度时进行聚合和更新,或者根据一定时间段内聚合到的信息进行更新。但目前ASP模式收敛性较差,与Synchronous模式相比还有很大差距。

 

采用更利于通信的拓扑结构

 

上文讨论了如何通过放宽同步要求来减少用于模型同步的通信,下面我们进一步分析如何更高效地支持这些通信,主要涉及图7中三种拓扑结构的讨论。

 

 

图7 三种通信拓扑结构[2]

 

 

1. Parameter Server

如图7(a),Parameter Server(PS)结构中有两种角色,Server和Worker。Server保存了模型的最新参数,并负责收集和处理不同Worker提供的梯度信息。每个Worker独立工作,负责载入各自的数据,从Server拉取最新的模型参数,并计算梯度等信息,发送给Server用于更新。在这一结构中,Server处于通信拓扑的中心,这是一种经典的中心化的框架,存在的主要问题是Server处容易出现通信拥塞(communication congestion)。

 

2. All-Reduce

由于PS结构中Server处容易形成通信瓶颈,大家更倾向于使用All-Reduce的结构进行梯度聚合,其结构如图7(b)。All-Reduce结构中仅存在一种角色,即Worker,每个Worker在通信过程中会得到其他所有Worker的梯度信息,用于完成本地模型的更新,所以在更新结束后所有Worker上的模型能够保持一致。这一结构能很好地支持Synchronous模式,非常适用于BSP-SGD,但对异步操作支持不佳。

 

3. Gossip

Gossip结构[8]如图7(c),这种结构中不仅没有一个中心的Server,也没有一个全局一致的模型,图中以不同的颜色表示不同Worker上维护一个不同的本地模型。在基于这一结构的训练过程中,每个Worker仅与neighbors(也叫peers)进行通信,每次迭代会选择一些特定的边,在这些边相连的两个Worker之间进行通信。在这种通信模式下,训练过程中不保证模型的全局一致性,仅在训练结束时给出全局一致的模型。基于PS结构采用SSP-SGD方法时与这一模式相似,但PS结构始终维护一个全局一致的模型,而Gossip不提供这一保证。

 

基于上述3种拓扑结构分别采用4种同步模式时,对训练过程的影响总结见表1,主要涉及4个方面的分析:模型一致性(Model Consistency)、通信频率(Communication Frequency)、通信拥塞情况(Communication Congestion)、模型收敛性(Convergence)。

 

 

表1 同步模式与拓扑结构的适配情况[2]

 

 

整体来看,保持模型的高度一致性可以使分布式训练的情况逻辑上更接近单机训练的情况,获得更好的收敛性(如表中PS+BSP和All-Reduce+BSP的例子),但高度一致性需要高通信频率来支持,在PS结构下会明显增加Server处的通信量,造成严重拥塞,而All-reduce结构可以大大缓解拥塞情况。ASP方式最大程度地放宽了模型一致性的要求,各个设备仍需保持高频通信,由于这些通信是异步进行,拥塞情况也大大减少,但很大程度上牺牲了模型的收敛性,且无法利用All-Reduce结构的优势。最后,正如上文所提到的,Local SGD方式在通信效率和模型收敛性上提供了一个可以权衡的选项。

 

总结与思考

 

面对各领域不断积累的大量数据,AI训练任务的规模迅速扩大,大规模训练的案例不断涌现,比如去年8月份微软发布的用512块A100训练的3.5万亿参数的MoE模型[11],再如去年双十一已经投入使用的在512块V100上训练的10万亿参数的M6-10T模型[12]。分布式训练系统中,不同的设备可能会运行模型的不同部分,或者处理不同的数据,而整个系统逻辑上是完成同一个模型的训练,这就需要设备之间通过大量的通信操作来维护全局一致性。由于目前互联带宽十分有限,这些通信操作很容易成为系统瓶颈。

本文从两个方面讨论了应对通信瓶颈的方法。一是放宽同步要求,减少用于同步的通信,主要思路包括允许较快的设备独立进行更多次迭代,使算力得到更充分的利用;增大通信间隔;使用异步策略。二是采用更利于通信的拓扑结构,主要思路为将通信分布到不同空间和不同时间,降低带宽压力。不同的同步模式配合不同的拓扑结构使用时会对训练过程产生不同的影响,需要使用者在模型收敛性和通信代价之间做出权衡。此外,还可从通信数据压缩以及通信过程和计算过程之间的调度两个方面来应对通信瓶颈,我们会在后续文章中进行介绍,欢迎大家持续关注。

 

参考文献

[1] AI and Compute. https://openai.com/blog/ai-and-compute/

[2] Tang, Zhenheng, et al. "Communication-efficient distributed deep learning: A comprehensive survey." arXiv preprint arXiv:2003.06307 (2020).

[3] Ruben Mayer, et al. "The tensorflow partitioning and scheduling problem: it's the critical path!." Proceedings of the 1st Workshop on Distributed Infrastructures for Deep Learning. 2017.

[4] Leslie G Valiant. "A bridging model for multi-core computing." European Symposium on Algorithms. Springer, Berlin, Heidelberg, 2008.

[5] Qirong Ho, et al. "More effective distributed ml via a stale synchronous parallel parameter server." Advances in neural information processing systems. 2013.

[6] Sebastian U Stich. "Local SGD converges fast and communicates little." arXiv preprint arXiv:1805.09767 (2018).

[7] Xiangru Lian, et al. "Asynchronous parallel stochastic gradient for nonconvex optimization." Advances in Neural Information Processing Systems 28 (2015): 2737-2745.

[8] David Kempe, et al. "Gossip-based computation of aggregate information." 44th Annual IEEE Symposium on Foundations of Computer Science, 2003. Proceedings.. IEEE, 2003.

[9] You, Yang, et al. "Reducing BERT pre-training time from 3 days to 76 minutes." arXiv preprint arXiv:1904.00962 (2019).

[10] Lin, Yujun, et al. "Deep gradient compression: Reducing the communication bandwidth for distributed training." arXiv preprint arXiv:1712.01887 (2017).

[11] DeepSpeed powers 8x larger MoE model training with high performance. https://www.microsoft.com/en-us/research/blog/deepspeed-powers-8x-larger-moe-model-training-with-high-performance/

[12] Lin, Junyang, et al. "M6-10t: A sharing-delinking paradigm for efficient multi-trillion parameter pretraining." arXiv preprint arXiv:2110.03888 (2021).

 

上一个: AI框架里的并行技术

下一个: 众包训练:另类的分布式异构深度模型训练方法

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多