当前位置: 首页 > 科技观察

八卦快手八卦:突破TensorFlow和PyTorch并行瓶颈的分布式训练框架

时间:2023-03-14 18:45:56 科技观察

近日,快手与苏黎世联邦理工学院宣布开源分布式训练框架八卦(Bagua)。学习开源框架只是在系统层面进行优化。八卦突破了这一点,针对分布式场景设计了特定的优化算法,实现了算法和系统层面的联合优化,性能较同类提升60%。研发背景随着摩尔定律的失效,单个计算单元的能力远远满足不了数据的指数级增长。例如,快手每天上传超过1000万个新视频。即使训练一个简单的分类模型(比如ResNet),以单机单卡的计算能力,一天之内训练完快手个新视频也需要100多天。.因此,在数据爆发式增长的互联网行业,多机多卡并行训练成为大数据时代的必然。随着深度学习模型的功能越来越强大,分布式训练任务的通信成本和所需的计算能力也急剧增加。但由于多机多卡并行带来的额外通信成本,提速(speedup)往往让大家大失所望,从而形成了大厂“堆资源”、无资源“发呆”的局面。例如,Google的Downpour框架[1]使用80个GPU来训练ImageNet,但加速比仅为12/80=15%。因此,如何提高多机多卡训练的通信效率成为并行训练乃至解决数据爆发式增长的核心问题之一。项目GitHub地址:https://github.com/BaguaSys/bagua现有的深度学习开源框架(PyTorch、TensorFlow)主要针对系统级优化,将现有的单机单卡优化算法扩展到多机机多卡场景。虽然系统层面的优化不断提升并行效率,但边际收益也越来越明显。针对这一问题,快手与苏黎世联邦理工学院联合开发了名为“八卦”的分布式训练框架,突破单纯的系统级优化,针对分布式场景设计特定的优化算法,实现算法和系统级的联合优化最大化分布式训练的效率。用户只需添加几行代码,即可将单机单卡训练扩展为多机多卡训练,并获得非常可观的加速比。八卦的设计思路从单机单卡训练到多机多卡训练。核心是每张卡累加并传播自己的计算结果。这个过程就像每个人把自己知道的信息传递给他人,再从别人那里获取信息,最后完成全球信息同步。如果将计算单元之间的信息同步比作人与人之间的信息同步,那么社会实践经验告诉我们,“八卦”可能是最高效的消息传递方式。“八卦”消息传播具有去中心化、异步传播、信息压缩等特点,与八卦中实现的传播算法恰好相呼应。为了提高分布式训练的效率,八卦实现了自主研发的前沿算法,包括去中心化/集中化、同步/异步、通信压缩等基础通信组件。这些组件的效率已通过软件和硬件的组合进行了优化。并灵活支持这些算法的组合,以及更复杂的算法设计。八卦将通信过程抽象为以下算法选项:CentralizedorDecentralized:在集中式通信模式下,梯度或模型的同步过程需要所有工作节点的参与,因此相对较高的网络延迟往往会导致训练效率降低。去中心化通信模式[5,6]往往可以有效解决这个问题:在这种模式下,工作节点可以连接成特定的拓扑结构(如环),在通信过程中,每个工作节点只与其自身通信邻居。SynchronousorAsynchronous:在同步模式下,在每次迭代过程中,所有工作节点都需要进行通信,下一次迭代必须等待本次迭代的通信完成后才能开始。相反,异步分布式算法[2]不需要等待时间:当一个节点完成计算时,它可以直接传递局部梯度来更新模型。全精度模式或信息压缩模式(Full-PrecisionorLow-Precision):在全精度模式下,将使用与本地模型相同的32位浮点数(float32)进行传输。另一方面,在通信瓶颈的情况下,基于大量现有研究,通过量化(quantization[3])或稀疏化(sparsification[4])等方法压缩梯度,然后将参数用压缩梯度更新。在很多场景下,它可以在提高通信效率的同时,达到和全精度一样的精度。虽然为了提高通信效率,八卦并没有像传统方式那样同步所有计算节点的结果,甚至每次同步的信息都存在偏差,但得益于最新的理论进展,这些通信策略及其组合最终会收敛。仍然可以充分保证解的正确性和效率,计算复杂度与同步集中和信息无损方法相当,但通信效率更高[10]。值得注意的是,在实践中,分布式训练算法往往使用以上一种以上的优化方法来适应更极端的网络环境[7,8,9]。对于对分布式算法感兴趣的读者,我们在这里推荐一份最近的完整调查报告[10]。八卦提供了一套详细的沟通模式,支持用户选择以上模式的任意组合。我们在下表中总结了本分布式训练系统对上述算法选项的支持情况:从表中不难看出,现有框架的优化只是针对更通用的算法(集中同步完全精度),针对其他算法组合,这些系统支持非常有限。对于集中同步的信息压缩,这些系统往往只支持相对简单的float32->float16压缩。相比之下,八卦可以支持更复杂的ByteGrad、QAdam等算法。对于其他的算法组合,现有的框架通常无法支持,但八卦可以自由支持。然而,简单地支持这个算法选项并不会直接导致在大型集群上的性能提升。八卦的核心优势在于算法和实现的联合优化,追求极致的性能。具体来说,基于上述对通信层的抽象,用户不仅可以方便地选择系统提供的各种算法组合来提升性能,还可以灵活实现一种新的分布式SGD算法——八卦将自动实现该算法提供系统层优化。这些系统优化包括:将通信时间隐藏在计算时间中:为了减少通信开销,八卦可以将部分通信时间隐藏在计算时间中。具体来说,在反向梯度的计算过程中,可以在剩余梯度的计算过程中同时通信部分完成的梯度——通过这种流水线处理方式,可以有效地将部分通信时间“隐藏”在计算过程中反向梯度,减少了数据并行带来的通信开销。l参数分桶和内存管理:频繁传输碎片化数据会降低通信效率。因此,八卦将模型参数划分成桶,分配连续的内存空间来管理每个桶,让通信的单位变成一个桶,从而可以更高效地使用通信模型。另外,由于支持信息压缩算法,压缩和解压功能的基本操作单位也是桶,也可以减少这些操作的开销。l分层通信实现:由于工业级分布式训练往往需要多机多卡,不同的物理连接方式带来的延迟和带宽也有较大差异,因此通信的有效抽象对性能也很重要。提升很关键。八卦将涉及多台机器的通信抽象为:“机内”和“机间”,并优化相应的通信抽象。例如,对于信息压缩传输,分级通信将这种算法解释为“机内”完全精度、“机间”信息压缩,从而为不同的物理链路提供最适合的通信算法。我们要强调的是,这些在系统实现层面的优化是广泛适用于各种算法组合的,而不是局限于特定的算法设置。因此,所有的系统优化都可以在各种算法实现中灵活复用,既保证了“端到端”的性能提升,也为开发新的分布式算法提供了良好的平台。经过实验,八卦的特点如下:l并行性能显着提升:在ImageNet数据集上,与目前开源的分布式框架(PyTorch-DDP、Horovod、BytePS)相比,在同等算力(128GPU)和通信网络(100Gbps),达到同样的训练精度,八卦只需要其他框架80%左右的时间;l对网络环境的鲁棒性更强:由于对各种算法优化(信息压缩、异步、去中心化)的有效支持,八卦在各种网络环境(包括不同时延和带宽)下表现出了良好的适应性。尤其是在高延迟、低带宽的情况下,八卦比其他框架表现出了更好的加速比。例如,在10Gbps的网络带宽环境下,八卦只需要其他框架50%左右的训练时间就可以实现相同的ImageNet任务。相同的训练精度;l“一键式”使用:八卦对终端用户非常友好。现有的使用PyTorch的模型都可以作为八卦的输入,八卦会自动为其提供丰富的并行解决方案——只需添加几行代码,训练就可以在分布式集群上运行;l分布式通信算法易于扩展:八卦为算法提供了高扩展性。对于分布式优化算法的开发者,八卦提供了有效的通信抽象。实现的新算法也可以直接复用八卦的系统优化;l可用于大规模工业场景:八卦为Kubernetes实现自定义算子,支持云原生部署,并考虑机器资源和故障问题,有机结合PyTorchElastic和Kubernetes实现容灾和动态训练扩容.用户可以在少量机器闲置时使用八卦开始训练,当更多机器资源释放时,训练任务自动扩展到更多机器上。同时,当机器节点损坏时,自动剔除坏节点继续训练。方便工业级训练场景使用,也方便结合机器学习平台使用;l安全易排错:八卦通信后端采用Rust语言实现,注重内存安全、速度和并发,编译时排除大量内存SecurityQuestion。同时,基于tracing实现了模块式、分级式的日志输出,在实际场景中更容易排查问题。此外,八卦还通过了快手内工业级任务的实战考验。八卦已经在快手内的多个核心业务场景中投入使用。与其他开源框架相比,取得了显着的性能提升:lLarge-scalenaturallanguageProcessing(GPT2-xlsizemodel),效率提升65%llarge-scaleimagerecognition(10亿+billionimages/videos),improve效率提升20%~30%l大规模语音识别(TB级语音数据),提升效率20%~30%l大规模推荐系统(万亿级参数模型训练,支持亿级应用DAU),效率提升100%以上八卦与其他开源方案的性能对比快手选择图像、文本、语音、图文混合等训练任务测试八卦的性能,并与PyTorch-DDP进行对比,霍洛沃德,BytePS。得益于八卦系统的高效性和算法的多样性,八卦可以针对不同的任务选择相应的最优算法,从而保证在训练精度与其他系统相当的前提下,训练速度得到显着提升。值得注意的是,当网络条件较差时,八卦系统的优势会更加明显。下面我们选取GPT2-XL、BERT-Large和VGG16三种典型的通信密集型任务进行对比说明。更多结果可以在八卦论文和网站(https://github.com/BaguaSys/bagua)中找到。1.端到端训练时间对比下图是在128个V100GPU上fine-tuneBERT-Large(SQuADdataset),模型的F1准确率随训练时间变化。八卦使用QAdam-1bit算法加速,机器之间通过100GbpsTCP/IP网络连接。我们可以看到,即使使用高速网络,八卦也只需要60%的时间就可以达到与其他系统相同的训练精度。2.可扩展性对比下图展示了各系统在VGG16模型(ImageNet数据集)上的训练速度与GPU数量的关系。测试中分别使用1、8、16、32、64、128个V100GPU进行测试。在这个任务中,八卦使用了8bitsGrad算法来加速。可以看出,八卦的扩展效率相比其他系统有明显的提升。下图展示了各系统在GPT2-XL模型上的训练速度与GPU数量的关系。GPT2-XL模型有15亿个参数,在大型模型中具有一定的代表性。八卦使用8bitsGrad算法加速。测试中分别使用8、16、32、64、80个V100GPU进行测试。也可以看出,八卦的扩容效率相比其他系统有明显的提升。3、不同网络环境的对比算法是八卦的灵魂。当网络环境发生变化时,不同的算法会表现出不同的性能特点。下图中,我们以BERT-Largefine-tune为例,调整机器间网络的带宽和延迟,并对比八卦中各算法的epochtime。可以看出,随着带宽的降低,压缩算法的优势会越来越明显,这与压缩的程度有关;当延迟逐渐增加时,分散算法将逐渐显示其优势。此外,当网络环境不好时,八卦相较于其他系统的优势进一步扩大。八卦使用示例在已有的训练脚本中使用八卦非常简单。在代码中,算法使用者只需添加如下几行代码即可对已有模型进行初始化。以GradientAllReduce算法为例:首先我们需要导入baguaimportbagua.torch_apiasbagua然后我们可以初始化八卦的进程组:torch.cuda.set_device(bagua.get_local_rank())bagua.init_process_group()for数据集初始化,八卦完全兼容PyTorch的实现:train_dataset=...test_dataset=...train_sampler=torch.utils.data.distributed.DistributedSampler(train_dataset,num_replicas=bagua.get_world_size(),rank=bagua.get_rank())train_loader=torch.utils.data.DataLoader(train_dataset,batch_size=batch_size,shuffle=(train_samplerisNone),sampler=train_sampler,)test_loader=torch.utils.data.DataLoader(test_dataset,...)最后用户只需要选择要训练的模型和优化器就可以使用八卦了:#定义模型model=...model=model.cuda()#定义优化器optimizer=...#选择八卦算法来使用frombagua.torch_api.algorithmsimportgradient_allreduce#实例化八卦算法algorithm=gradient_allreduce.GradientAllReduceAlgorithm()#对已有模型启用八卦算法model=model.with_bagua([optimizer],algorithm)这样,使用八卦的多机多卡训练算法就完成了。完整示例和更多场景请参考八卦教程文档(https://github.com/BaguaSys/bagua)。可用:https://arxiv.org/abs/2107.01499开源GitHub托管:https://github.com/BaguaSys/baguaInstallationRepository[1]Dean,Jeffrey,GregS.Corrado,RajatMonga,KaiChen,MatthieuDevin,QuocV.Le,MarkZ.Mao等人。“大规模分布式深度网络。”(2012).[2]ZhengyuanZhou、PanayotisMertikopoulos、NicholasBambos、PeterGlynn、YinyuYe、Li-JiaLi和LiFei-Fei。2018.具有无限延迟的分布式异步优化:你能走多慢?。在机器学习国际会议上。PMLR,5970–5979.[3]DanAlistarh、DemjanGrubic、JerryLi、RyotaTomioka和MilanVojnovic。2016.QSGD:通过梯度量化和编码实现通信高效的SGD。arXiv预印本arXiv:1610.02132(2016).[4]DanAlistair、TorstenHoefler、MichaelJohansson、SaritKhirirat、NikolaKonstantinov和CedricRenggli。2018.稀疏梯度方法的收敛。在第32届神经信息学国际会议论文集中在处理系统上。5977–5987.[5]AnastasiaKoloskova、SebastianStich和MartinJaggi。2019.具有压缩通信的分散随机优化和八卦算法。在机器学习国际会议上。PMLR,3478–3487.[6]XiangruLian、CeZhang、HuanZhang、Cho-JuiHsieh、WeiZhang和JiLiu。2017.去中心化算法能否胜过中心化算法?分散并行随机梯度下降的案例研究。在第31届国际神经信息处理系统会议论文集中。5336–5346.[7]ChristopherDeSa、MatthewFeldman、ChristopherRé和KunleOlukotun。2017.理解和优化异步低精度随机梯度下降。在第44届年度国际计算机体系结构研讨会论文集中。561–574.[8]连相如、张伟、张策、季六。2018.异步分散并行随机梯度下降。在机器学习国际会议上。PMLR,3043–3052.[9]唐汉林、干少多、张策、张彤、刘季。2018.分散训练的通信压缩。在第32届神经信息处理系统国际会议论文集中。7663–7673.[10]刘季,张策,等。2020.具有一阶方法的分布式学习系统。FoundationsandTrends?inDatabases9,1(2020),1–100。