3.1. 设计综述

3.1.1. 背景

纵观深度学习的发展史,不难发现,很多奠基性的工作,其实早在上世纪四五十年代就被提出了。但是囿于算力,当时的成果主要集中在理论上。直到最近十年间,随着计算性能的不断提升,我们能够在有限的时间内训练出更大更深的神经网络,才让深度学习得以腾飞。不夸张的说,深度学习如今在各个领域取得的成就,和大规模是分不开的。

不管是学术界还是工业界,都一直致力于更快速的训练更大更深的神经网络。大体来分,我们在深度学习领域主要有两个追求,一是训练性能更快,二是模型规模更大。分布式深度学习领域中的很多概念和技巧,都是为了解决这两个问题产生的。

接下来我们着眼于分布式训练从性能优化和大模型训练两方面入手,介绍几个常用的方法。

3.1.2. 性能优化

在通用GPU发布之后,使用显卡训练神经网络的热度开始爆炸性地增长。NVIDIA的CUDA编程语言可以让用户以一种像C一样的语言实现任意代码。那么如何在通用GPU上设计出高效的代码,对性能提升就变得至关重要。本小节大部分的内容,正是关注于此。值得一提的是,除了GPU,市场上还涌现了很多其他硬件厂商开发的AI专用芯片,例如百度的昆仑、华为的昇腾910等。当然,在这些不同芯片上的优化思路都是类似的。

在分布式机器学习中,最常用的并行模式是数据并行,即每个工作节点拥有全部模型参数,并训练全量数据的一部分,之后对计算出来的梯度(或参数)进行通信(通常为all-reduce操作)以实现全局信息共享。可以看出,计算和通信是分布式深度学习任务中最主要的两部分。所以常用的性能优化策略也正是从这两部分入手进行的。在计算方面优化手段主要包括计算算子融合;在通信方面优化手段主要包括通信算子融合、通信拓扑优化等;当然,还有一部分优化同时涉及两部分,例如混合精度训练。下面将会逐一介绍。

3.1.2.1. 计算OP融合

在深度学习框架中,最基本的计算单元是算子(Operator)。例如常见的矩阵乘法操作,就是以MatMul算子的形式存在。一个完整的计算网络,通常就是由多个算子组合起来的。这样的设计十分灵活,用户可以通过组合不同的算子来验证不同的想法。

但是,鱼和熊掌不可兼得。拥有巨大灵活性所要付出的代价就是性能。举例来讲,假设我们要计算三个输入a、b、c相加的结果,调用过程可能是tmp=add(a, b); out=add(tmp, c)。在这样的网络中,我们会启动两次计算,并开辟了一个中间变量用于存放中间计算结果。在CUDA开发中,这样的一次计算通常是由一个或多个Kernel进行的,而Kernel的启动通常需要一定时间开销,因此在这个例子中加法运算所用到的Kernel就要启动两次,即将产生双倍的时间开销。

针对这个操作的一种优化方法是,我们开发一个支持三个输入的OP(假设名为add3)。那么我们只需要启动一次Kernel计算,即out=add3(a, b, c),便可以得到最终的结果。该方法的一个附加好处是还节省了一个临时空间的申请。

这种思路就是所谓的计算OP融合(Fusion),详细内容请参考计算融合。需要说明的是,OP融合在单卡下就有效果,并不是分布式特有的策略。对分布式训练来讲,如何在计算和通信并重的情况下获得更优秀的性能,是我们关注的重点。

接下来的几个小节会结合一个生动的例子来阐述各种优化策略的思想。我们的主人公是Alice和Bob两位小朋友,他们要在各自的房间里做一沓试卷,每张试卷上有若干题目,覆盖不同的知识点。他们的目标是做完所有的试卷,并学到相应的知识。特别的,他们可以通过交换各自学到的内容来修正或巩固自己的知识。Alice和Bob一开始选定的做法是:每当他们之间有人做完一道题,就拨电话给对方,等对方也做完这道题并接起电话后,同步各自的答案,然后同时开始做下一道题。

3.1.2.2. 通信OP融合

Alice和Bob所在的国家电话号码很长,所以他们发现每做完一道题就互相通话,拨电话号码的耗时有些难以接受。他们想,如果商定好做完多道题目,再通话一次进行交流,能省去很多拨电话号码造成的时间开销。

这就是通信OP融合的思想。我们知道每次触发通信都会有一些额外的操作(如先建立连接等),减少这些额外的操作将对性能有很大帮助[1]。顺着这个思路,如果我们能够将多次通信的内容先拼接成连续的数据,然后在一次通信内全部发送/接收,那么将会更充分的利用硬件资源,获得更大的性能提升。

通信OP融合的使用方法请参考通信融合

3.1.2.3. 计算和通信重叠

按照之前的约定,做题快的人(比如Alice)拨通电话后,要等待Bob完成对应的题目之后接起电话才能开始这次通信。在等待Bob接听电话的时候,Alice只是闲坐在那里听着听筒里的彩铃音乐。她突然想到,为什么要听这种无聊的声音,而不开始提前做下面的题目呢?

这就是通信和计算重叠的思想。CUDA中有流(stream[2])的概念,一个流表示一个GPU操作队列,该队列中的操作将以添加到流中的先后顺序而依次执行。那么通过令计算和通信操作加入不同的流中,可以做到二者的执行在时间上重叠。详细内容请参考通信重叠

3.1.2.4. 通信拓扑优化

现在做题的团队壮大了,除了Alice和Bob,又加入了几位新同学。他们的目标变成要让每个人算出来的答案,都被所有其他人知道。最简单的做法,自然是所有人之间通一次电话。但是这样做时间开销太大了。聪明的他们选择了另一种做法,把所有人分成几组,每个组选出一名组长,组员把答案汇总给组长。组长间先互相交换所有的信息,然后再分发给所有组员。

不同的信息交换策略,对应到分布式训练中,就是不同的通信拓扑。上述采用的通信策略借鉴了分层(hierarchical)通信的思想。在业界,有ring-allreduce[3],Double binary trees[4]等多种拓扑结构。

通信拓扑优化的更多使用方法,请参考通信拓扑优化

3.1.2.5. 深度梯度压缩

再次回到仅有Alice和Bob两人做题学习的场景来。他们在做题过程中发现,随着学习的进行,对于不同知识点的掌握程度有好有坏。有的知识点已经掌握的很好了,再做题也提供不了太多新的知识。但另外一些,却仍然感到模棱两可。于是两人约定,每做完T张试卷,选出最拿不准的几个知识点来交流答案,而掌握充分的那些知识点,就不在电话中交流了。

上述思路就是深度梯度压缩(Deep Gradient Compression, DGC)的主要思想。DGC通过将梯度稀疏化,在每轮训练时只选择出一部分比较“重要”的梯度进行同步,以达到降低通信量的目的。当然,减少通信量势必会造成精度损失。为了减少损失程度,作者还提出了动量修正(momentum correction)、本地梯度裁剪(local gradient cliping)、动量因子遮蔽(Momentum factor masking) 等几项技巧。详细内容可以参考DGC优化低配网络的分布式GPU训练

3.1.2.6. Local SGD

Alice和Bob觉得没必要每道题都打电话交流答案,就算使用了前述通信OP融合的技术,也只是减少了打电话的频率,但还是每一道题都要对答案。

于是两人又想到了一个能够减少打电话次数的策略:他们决定各自先做T张试卷,自行学习梳理各个知识点的知识,然后再通电话交流各个知识点的心得。当按照这个方法执行的时候两人发现,尽管花费在打电话上的时间确实减少了,但副作用是他们各自学到的知识可能不一定准确,交流次数的减少让他们没法及时纠正自己某些错误的理解。因此他们又想到了另一个更好的沟通策略:刚开始学习的时候交流频繁一点,当对各个知识点有了大致的了解后,再慢慢降低通话的频率。毕竟具备了基础知识后,只有在题海中遇到新题才能带来新的认识,刷再多重复的题目是没什么意义的。

Local SGD就是基于这个思路,最基本的Local SGD属于上例的第一个策略,直接增大参数同步的间隔来减少通信耗时,但是弊端是可能造成训练精度的损失。而对于第二种策略,Local SGD又衍生出了post Local SGD和Adaptive Local SGD两款“加强版”:

  • post Local SGD训练的第一个阶段保持每算出一个参数的梯度,就完成一次同步通信,以保证训练精度;之后到了第二阶段,则增大同步间隔(该间隔是固定的),以提升训练效率。

  • Adaptive Local SGD相对于post Local SGD而言,更加灵活,它会动态调整梯度同步通信的间隔,从而达到训练精度和训练速度之间的平衡。

详细内容可以参考使用Local SGD优化低带宽下分布式训练

3.1.2.7. 自动混合精度

Alice和Bob有一个特殊记忆能力,就是可以把想表述的内容,提炼成少量文字(原先的字数的一半),这样可以减少记忆的内容,但是同时也会导致准确性稍稍出现偏差。

随着知识点和题目越来越多,Alice和Bob觉得脑子发沉,可能脑容量已经快用完了,而打电话交流的时间也越来越长。于是他们决定用上那种记忆能力,这样就释放了大脑中更多的空间,而且打电话交流的内容也随之减半。

在实际应用中,对应这种记忆能力的就是半精度(FP16)类型,使用半精度类型进行训练,称之为混合精度训练(AMP)。混合精度训练有若干好处,例如减小显存使用量,增大通信吞吐等。当然精度的降低会导致数字表示范围的缩小,进而导致比FP32更容易溢出,为了应对这些问题,我们引入了Dynamic loss scaling和op黑白名单等策略来避免。

  • Dynamic loss scaling:在AMP训练过程中,为了避免精度下溢,每训练一定数量批次的数据,就将Loss放大指定倍数。如果Loss在放大过程中发生上溢,则可以再缩小一定倍数,确保整个训练过程中,梯度可以正常收敛。

  • op黑白名单:通过使用大量模型在不同应用场景中反复验证后,飞桨团队根据半精度数据类型计算的稳定性和加速效果,梳理出一系列适合转换为半精度计算的算子,并将这些算子定义到了一份白名单文件中。同时对于一些经过验证发现不适合转换的算子,也就是使用半精度计算会导致数值不精确的算子将被记录到黑名单文件中。此外一些对半精度计算没有多少影响的算子归类于灰名单。在使用自动混合精度训练过程中,系统会自动读取黑白名单,从而感知到哪些算子需要被转换为半精度计算。

详细内容请参考自动混合精度训练