
还是在今年(2018年)11月美国感恩节放假期间,我无意点开我的新论文搜索关注,假期的懈怠顿时被一扫而空。一篇谷歌的新论文跳入眼帘,声称打破了几天前刚建立的分布式训练速度的记录。各大公司训练速度记录上的竞赛早已白热化了。11月13日,Sony宣布,动用1088台服务器,4352卡最强大的Nvidia Tesla V100的集群,创造了3分44秒训练ResNet50的记录。仅仅3天后,谷歌宣布,在几个月前刚发布的1024片TPUv3助力下,该记录被缩短到2分12秒。这场竞赛怎么开始的?各大世界级技术公司和实验室如何不惜血本刷记录的?拼比哪些关键技术呢?
相关服务:日本站群服务器
知名团队如同世界杯百米赛般刷记录固然是引人入胜的故事,但是对我们更重要的是,分析分布式训练的主要技术方向,总结现有的经验教训。
本文先从深度学习的训练基础数学入手,分析出计算特点,影响体验的主要指标,然后从算法到系统各方面常用加速训练手段;通过总结各家团队刷记录的系统和技术,可以看到如何应用这些手段,以及进一步的思考。
数学基础和定义
先简单解释深度学习的训练过程,帮助说明一些后面要用的概念,如minibatch大小,learning rate和损失函数等高线等等。具体的公式和更详细的理论请参考【1】。
训练数据向量

反向传播时的损失函数由以下公式,

最常用的模型训练优化方法是梯度下降,Stochastic Gradient Descent,通过反向传播损失函数的梯度,如下公式,修正模型权重向量

迭代中计算梯度下降时,为了提高并行度,更新权重时采用一个批次梯度平均值,叫Batch Gradient Descent 或Minibatch Gradient Decient。定义公式如下,

大量的训练样本迭代,逐步修正模型向量序列,最终找到最佳解,使得损失函数最小。或者说,多次样本产生梯度下降,在模型向量序列的解法空间收敛到最佳点。
如果将模型向量序列想象成2维,解法空间可以由2维轴表达,每个点是一组解,这个解的损失函数的等高线图可以在上面画出。如红线所示,从起始点(解)开始,每次梯度反向传播修正模型,得到一个新解。经过若干次的梯度反向传播和修正,最终找到解法空间中损失函数的最小解。

训练计算和系统的特点
训练的计算与传统意义上的计算有很大不同,差异总结如下。

传统计算执行一系列有确定次序的操作和计算,最终得到的结果有逻辑上的确定性,不包含任何随机成份。每一步操作必须结果正确,否则误差会保留在一系列的操作后,影响最终结果。
训练的实质是迭代搜索全局最优解。训练成功与否衡量标准是能否找到损失函数最小点,但最优点实际中往往是相对的,不是唯一的,具备一定的不确定性。例如,如损失函数的等高线图,全局最优点为五角星处。搜索中也可能收敛到一个局部最优点,如红色的三角形处,这种情况被称为模型泛化不佳(generalized poorly);也可能收敛到真正最优点附近的某点,损失函数足够小到可以接受。是否收敛到理论最优点实际中较难判断,通常用参照模型准确度来衡量。譬如各公司刷新训练ResNet50的记录时,以已公开的ResNet50模型的准确度为参照。【2】中使用的模型参照Top-1 validation error 23.6%( 精确度76.4%);【14】使用Top-5 为参照93%。甚至模型准确度成为训练速度,算法和系统等等诸多因素中寻求平衡的一个维度,如【7】【8】【9】中讨论。

训练是迭代收敛过程,带来的一个好处是训练系统的容错性。迭代收敛过程能自动恢复小计算误差的影响,不一定会妨碍找到最优解【19】。如下图示意,图中描述在解法空间迭代求解,从第0步开始,迭代k步,目标是找到最优解

训练的迭代过程有相当的灵活度。只要训练能收敛到最优解,训练的迭代求解过程可以走多种多样的路径,不必是唯一的。迭代次序也是可以变化,不是唯一的。参数和算法有选择空间,可以调整,调整后迭代求解的过程有相应的变化。 如果能考虑各个算法和参数,找到合适的平衡点,训练仍然可以成功收敛。如下图示意,优化算法可以修改。不同的优化算法,如SGD,Momentum, NAG,Adagrad,Adadelta,Rmsprop,会迭代经过不同的路径,但都可以收敛到最优解。

除了优化算法,调整其他参数或选择不同算法,如minibatch大小, learning rate,weight decaying, 梯度同步方法,模型同步方法,梯度精度,模型精度,都会改变迭代训练的过程。限于篇幅,关于这些参数详细说明请参考【1】。minibatch大小和learning rate,已经在前文介绍过。后文中各团队刷新训练速度记录的技巧总结里可以看出,这两个参数对训练收敛非常重要。
从系统设计角度看,这些参数或算法选择实际上决定了分布式训练系统的设计。譬如,minibatch大小决定分布式训练系统可以用多少个节点平行前向传播和反向计算梯度;模型精度决定分布式训练系统的计算精度,从而决定芯片计算复杂度和需要内存大小;梯度精度和梯度同步方法要求各节点之间任务同步,连接,互联带宽和延迟需求。这个训练计算的特点为分布式训练系统设计提供灵活的设计空间,设计需要在各个因素或选择中间寻找平衡点,完全不同于传统计算系统设计。系统可以针对某些瓶颈优化,如增大batch大小/并行度来加快训练,然后调整算法或参数,保持模型训练的精确度和控制系统的产品质量。后文中各团队刷新记录过程一次又一次地重现这一点。因此,系统和算法co-design不会是选项,而是必然。
当考虑传统计算产品时,Total Ownership Cost(TCO)和计算速度是最重要的考量。训练计算产品的TCO当然同样重要,但用户关最终心的速度不再单纯是计算速度,而是模型训练收敛的速度。这两个并不等同。除计算能力,训练系统上可以包括其他设计,较大程度得影响收敛速度。由于训练出来的模型精确度会因为算法和系统不同而不同,并且有可能作为平衡训练速度和系统开销的因素之一。
总结一下,训练计算和训练系统有如下特点。
- 训练计算是以搜索最优解为目的,迭代收敛的过程,不同与传统计算以操作执行,严格次序的逻辑过程。
- 迭代收敛的计算过程能容忍一定的误差或错误。
- 训练中的算法和参数选择有一定灵活度,迭代只需要遵从松懈的次序。
- 训练出来的模型精确度是产品评估和设计指标之一。
- 训练计算和传统计算都需要考虑运算速度,但训练计算关心的是收敛速度。
- 训练计算和传统计算都需要考虑Total Ownership Cost。
- 训练的算法和系统因素可以相互弥补, 必须co-design。
大集群训练记录属于财大气粗的
Facebook最早提出大Batch大规模并行高速训练,引起业界的注意,带动一系列的记录刷新。各公司更关注较成熟和应用较多CNN的训练,基本以ResNet50在ImageNet上训练速度为指标,有几个尝试也训练了AlexNet或Inception的某个版本。原始ResNet50的参照是模型训练90个epoch的数据后,经validation set测试,达到76%Top-1,93%Top-5的精确度。我们看看训练计算的特点如何在这些工作中体现出来。
2017年7月,Facebook脸书Goya Welcome!. 发表了论文Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour【2】,提出采用大batch,大机群快速训练,使用ImageNet数据集1小时内训练ResNet50到76.3% Top-1精确度,minibatch大小8192,分布式系统包括32台Facebook自己设计的Big Basin服务器,网络为50Gb Ethernet;每台Big Basin插有8张Nvidia P100 GPU,卡和卡之间通过NVLink互联;软件使用Caffe2,GPU互联使用NCCL。论文提出算法技巧,特别强调了linear scaling rule 和新的warmup策略。 1)Linear Scaling Rule指按minibatch大小线性比例设置。 2)(slow rate warm up)开始训练时,为了减小优化难度,用小的learning rate。 3) 如果使用Momentum SGD,learning rate更新时需要做调整。 4) 按整个minibatch正则化损失函数计算。5) 一次训练数据随机换序应该分给所有使用同一epoch的节点。Gradient aggregation策略 3个阶段。 1)同一服务器上GPU的梯度累加求和; 2)所有服务器梯度和再累加成为总和;3)总和广播给回所有GPU。两个优化算法用于减少服务器之间的通信,recursive halving and doubling 算法【3】【4】和 bucket算法【5】(详细算法参见reference)。
2017年9月, UCBerkeley You et.al. 发表了论文ImageNet Training in Minutes【6】。11月, 日本的Preferred Network,Akiba Welcome!. 发表了论文 Extremely large minibatch sgd: Training resnet-50 on imagenet in 15 minutes【7】. 两篇论文增加Minibatch 大小到32k,ResNet50在ImageNet上的精确度达到74.9%Top-1所需训练时间减少到14 或 15分钟。他们使用了不同的系统和方法。Preferred Networks的系统使用128台服务器;每台两个Intel Xeon E5-2667 processors (3.20 GHz, eight cores), 8卡Tesla P100 GPUs (共1024卡);服务器之间用Mellanox Infiniband FDR连接。论文中通信没有明确提到NVLink,但说了使用NCCL(通常在NVLink上),梯度精度减少到16bit的half-float。软件用了Chainer。算法上延用Facebook的技巧。另外为了改善开始训练时优化的困难,开始时优化用RMSprop,慢慢转换到SGD;同时learning rate schedule从更小的初始值开始得更慢。另一个技巧是Batch Normalization不再用moving average。UCBerkeley的分布式系统包括2048个 Intel KNL CPU(Intel Xeon Phi Processor 7250),软件基于Caffe,MPI,Intel Machine Learning Scaling Library。算法上提出了更复杂的基于每层权重和梯度的learning rate机制,LARS (Layer-wise Adaptive Rate Scaling )【8】。但是这两项工作的模型精确度有损失,不能到达76%。UCBerkeley的结果稍好些,能在20分钟内训练完90个epoch到75.3%。Preferred Network的论文中没有提供训练数据量。
11月谷歌Smith Welcome!. 也发表了他们的工作, Don't Decay the Learning Rate, Increase the Batch Size 【9】,在30分钟内训练ResNet50到76.1%Top-1。使用了整个TPU Pod,共256个TPUv2。论文中没有给出详细的系统设置或节点通信机制。Cloud TPU Pod将所有TPUv2通过高速互联联成2维环形网格,通信带宽应该相当可观。在下面Ying et.al.中讨论了TPUv3利用2维互联做gradient aggregation,TPUv2应该有同样的能力,采用同样的方法。Cloud TPU中的向量计算使用谷歌自定义的bfloat16精度,因此是混合精度的。 算法改变以前常用的训练过程中逐渐减小learning rate的做法,而是逐渐增加minibatch大小。实验中前30个epoch将minibatch大小定为8k,然后增大到16k。


腾讯Jia Welcome!. 2018年7月再次刷新了记录,发表了Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes【10】。论文标题中的4分钟指的是训练AlexNet, ResNet50训练到75.8%Top-1精确度需要6.6分钟。该工作采用minibatch大到64k,集群包括256节点,每节点8卡Tesla P40 GPU,共2048卡。节点之间通过Mellanox ConnectX-4 100Gbit Ethernet连接,RDMA被用来并行计算和通信。软件使用了Tensorflow,NCCL和OpenMPI。算法上采用3个技巧:UCBerkeley You Welcome!. 提出的LARS,混合精度(fp16和fp32)计算, 停用了模型的bias上和batch normalization参数上的weight decay 。该工作提出了新的梯度通信策略:1) Tensor Fusion,为充分利用网络带宽,将几个小梯度tensor合在一起传送;2)Hierarchical All-reduce,如图,将每k个GPUs分一组, gradient aggregation的allreduce分成3个阶段。第一阶段组内的GPU做ring based allreduce;第二阶段通过RDMA在组和组之间allreduce;最后每组将结果广播给组内GPU。Gradient Aggregation时采用half float精度。

2018年11月13日,Sony的Mikami Welcome!.,发表了ImageNet/ResNet-50 Training in 224 Seconds【11】,宣布训练ResNet50到75.03%Top-1精确度只需要3分44秒。系统使用1088个节点,每节点上4卡Tesla V100用NVLink互联,2 Xeon Gold 6148 CPU, and 376 GB RAM。节点之间通过Infiniband EDR x2互联。软件使用Sony自研并公开代码的Neural Network Libraries (NNL), NCCL version 2.3.5 and OpenMPI version 2.1.3。 算法上采用上面已经提过的LARS, slow warmup learning rate, momentum SGD,不带moving averaging的batch normalization,minibatch大小管理(从32k增加到64k), 通过NCCL实现2维网状环形allreduce。注意同节点GPU之间连接通过NVLink,不同节点GPU必须通过Ethernet,所以2维网状环形不同连接的带宽不同。文中特别指出,尽管整个集群有1088节点,4352卡V100,当使用超过2176卡V100时,加速效率大幅度下降,所以最终创造记录时实际使用了2176卡。
记录只保持了3天,最新的刷新来自谷歌。Ying et.al.在2018年11月16日发表了Image Classification at Supercomputer Scale 【12】。新记录是ResNet50训练到76.3%Top-1精确度需要2.2分钟。训练系统用了一台今年谷歌新发布的TPUv3 Pod,共1024 TPUv3。 TPUv3和TPUv2一样,采用高速2维环状网络芯片互联,计算使用混合精度,但每片计算能力是v2的2.3倍。而且TPUv3的系统功耗更小,机板更紧凑,如图,每TPUv3 Pod机架可容纳4倍的TPU机板。TPUv2的CPU服务器需要单独占一个机架的,TPUv3的CPU服务器机板可以和TPU机板合放在同一机架。论文中的算法技巧包括,1)控制learning rate时,使用了linear learning rate scaling,缓慢增加然后减小learning rate(slow warm-up and decay); 2)You Welcome!. 的LARS; 3)大规模分布式训练,为减少节点之间通信,往往让每节点各自做Batch大小为32的Batch Normalization。论文提出新的distributed batch normalization方法。如图,若干节点分为一组,先在本节点计算mean和variance,每组计算本组的mean和variance,再送回各节点更新tensor。
论文提及不少系统上的优化。为了避免输入数据处理成为瓶颈,系统数据流程中主要优化了数据集缓存,提前读取,图像解码和剪裁,数据解析。Gradient aggregation充分利用TPUv3强大的2维环状网络芯片互联。如图,用3x3个TPUv3的互联示意。每个TPUv3有两个core,图中分画为红色和蓝色。第一步,蓝色的core沿垂直方向累加和传递梯度值,同时红色的core沿水平方向累加和传递梯度值;第二步,蓝色的core沿水平方向累加和传递梯度值,同时红色的core沿垂直方向累加和传递梯度值。



平民化(一些)的训练尝试
以上快速训练的尝试假设计算量固定为ResNet50的基准90 epoch,依靠加大minibatch,大并行集群和新硬件的策略。2018年谷歌之外的云供应商推行一个新的策略。该策略依赖AWS或云现有集群,依赖现有NVLink的NCCL的ring allreduce,通过改进训练的方法,算法和软件,试图加快训练速度。虽然训练时间离大minibatch,大并行集群和新硬件的策略还有差距,但是经验值得借鉴。需要注意的是,因为改进训练算法的目的是为了迭代收敛得更快,训练所需的epoch数会小于原始ResNet50训练的基准90 epoch,所以两策略的计算量是不同的。
开始这一方向的是Amazon AWS。2018年7月,AWS在blog上发表了题为Scalable multi-node deep learning training using GPUs in the AWS Cloud【13】 的博文。文中宣布在8台AWS P3.16xlarge instances 上47分钟内训练ResNet50到75.75%Top-1。一个P3.16xlarge instances上包括64 vCPUs (定制的Intel Xeon E5), 488GB RAM,25Gbps的网速,8卡Tesla V100 GPU。8 instance总共64卡V100 GPU【16】。使用ApacheMXNet时,ResNet50训练47分钟完成90个epoch, 到达75.75%Top-1。使用TensorFlow+Horovod时,50分钟到75.54%。算法上主要使用linear rate scaling,slow warm up和fp16/fp32混合精度。
2018年8月,fast.ai · Making neural nets uncool again 发表了题为Now anyone can train Imagenet in 18 minutes【14】的博文。原始的ResNet50的基准Top-5精确度是92.2%,博文中宣布在16台AWS P3.16xlarge instances (共128卡 V100)上18分钟(6秒)内训练ResNet50到93%Top-5。这个数据也公布在Stanford的DAWNBench【15】。有趣的是,fast.ai在DAWNBench还公布了使用8台AWS P3.16xlarge instances(共64卡V100)训练ResNet50,只需要18分钟53秒。就是说,使用两倍的硬件资源,额外64卡V100,在训练速度上只从18分53秒加速到18分6秒,加速了47秒,说明在这样系统规模和这种训练速度,并行上的通信和管理开销已经是最主要的瓶颈,腾讯【10】和谷歌【11】论文中强调系统和通信的优化也印证了这一点。软件上使用了Pytorch和fast.ai自制工具。算法上的技巧上包括:1)progressive resizing, 开始训练时用小图像,训练过程中逐渐增加图像尺寸;2)直接用原始长方形的训练图像,传统流程将图像剪裁到正方形; 3)如【10】,停用batch normalization里面的weight decay; 4)dynamic batch sizes 【17】;5) slowl learning rate warm-up【2】。 优化算法使用了momentum SGD。
华为云在2018年10月的华为全联接大会发布了AI开发平台 ModelArts,特意提及用ModelArts在AWS上同样硬件训练ResNet50比fast.ai快。11月底在DAWBench提交了10分28秒训练ResNet50到93.16%Top-5精确度的记录。华为云公众号的发布文宣【18】里还宣布使用128各GPU 10分6秒可达到75%Top-1, 并且提到以下4个改进: 1)数据读取和预处理中采用并发输入流水线;2)采用混合精度计算和自适应的尺度缩放减小由于精度计算带来的损失;3)采用动态超参(如momentum、batch size等)减小所需epoch个数;4)服务器和通信计算库有所优化。
思考和展望
ResNet50在2015年发表时【20】在ImageNet数据集训练完成需要29小时, 经过各公司和实验室的这一系列的记录刷新,缩短到今天的2分12秒。这无疑是惊人的进步,同时提供了许多的技术和经验。然而早期学术界对分布式机器学习系统的研究曾经探索过更多方法【21】【22】【23】, 前文列出的训练算法和参数清单上许多手段并未充分尝试。分布式训练还有更长的路。
- 训练技术上还有潜力可以挖掘。
- 新的硬件和系统可以提供更强大的工具。
- 记录刷新研究集中在ResNet50 和少数CNN网络。其中有些技术并不能普适其他网络,甚至局限于ResNet50【10】。
- 记录刷新研究只考虑了数据并行的任务分配方法,模型并行并不在考虑范围。
- 其他网络和应用,如LSTM训练,Reinforcement Learning和Neural Architecture Search,会有不同的计算和系统需求。
【1】Good Fellow et. al., 2016, Deep Learning, MIT Press, http://www.deeplearningbook.org/
【2】Goyal et. al., 6/2017 1st ver., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. https://arxiv.org/pdf/1706.02677.pdf
【3】R. Rabenseifner. Optimization of collective reduction operations. In ICCS. Springer, 2004. https://pdfs.semanticscholar.org/8d44/e92b3597d9e3f5245e152c9e0ce55b3e68a4.pdf
【4】R. Thakur, R. Rabenseifner, and W. Gropp. Optimization of collective comm. operations in MPICH. IJHPCA, 2005. https://www.mcs.anl.gov/~thakur/papers/ijhpca-coll.pdf
【5】M. Barnett, L. Shuler, R. van De Geijn, S. Gupta, D. G. Payne, and J. Watts. Interprocessor collective communication library (intercom). In Scalable High-Performance Computing Conference, 1994.
http://www.netlib.org/intercom/shpcc94.ps
【6】You et. al., 9/2017 1st ver., ImageNet Training in Minutes, https://arxiv.org/pdf/1709.05011.pdf
【7】Akiba et.al.,11/2017 1st ver., Extremely large minibatch sgd: Training resnet-50 on imagenet in 15 minutes. https://arxiv.org/pdf/1711.04325.pdf
【8】You et. al., 9/2017 1st ver., Scaling SGD Batch Size to 32K for ImageNet Training, https://www2.eecs.berkeley.edu/Pubs/TechRpts/2017/EECS-2017-156.pdf
【9】Smith et. al., 11/2017 1st ver., Don’t decay the learning rate, increase the batch size, https://arxiv.org/pdf/1711.00489.pdf
【10】Jia et. al., 7/2018 1st ver., Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes, https://arxiv.org/pdf/1807.11205.pdf
【11】Mkami et.al., 11/2018 1st ver., ImageNet/ResNet-50 Training in 224 Seconds, https://arxiv.org/pdf/1811.05233.pdf
【12】Ying et. al., 11/2018 1st ver.,Image Classification at Supercomputer Scale, https://arxiv.org/pdf/1811.06992.pdf
【13】Ragab et. al., 7/2018, Scalable multi-node deep learning training using GPUs in the AWS Cloud, https://aws.amazon.com/blogs/machine-learning/scalable-multi-node-deep-learning-training-using-gpus-in-the-aws-cloud/
【14】Fast.ai,8/2018, Now anyone can train Imagenet in 18 minutes, https://www.fast.ai/2018/08/10/fastai-diu-imagenet/
【15】Stanford DAWN Project,12/15/2018的内容, DAWNBench An End-to-End Deep Learning Benchmark and Competition, https://dawn.cs.stanford.edu/benchmark/ImageNet/train.html
【16】Amazon AWS,Introducing Amazon EC2 P3 Instances,https://aws.amazon.com/about-aws/whats-new/2017/10/introducing-amazon-ec2-p3-instances/
【17】Friedlander et.al., Hybrid deterministic-stochastic methods for data fitting. SIAM Journal on Scientific Computing, 34(3):A1380– A1405, 2012.
【18】华为云, 华为云ModelArts深度学习训练全球最快,华为云公众号
【19】Qiao et.al., 10/2018 1st ver., Fault Tolerance in Iterative-Convergent Machine Learning, https://arxiv.org/pdf/1810.07354.pdf
【20】He et.al., 12/2015 1st ver., Deep Residual Learning for Image Recognition, https://arxiv.org/pdf/1512.03385.pdf
【21】Xing et.al., 12/2013 1st ver., Petuum: A New Platform for Distributed Machine Learning on Big Data, https://arxiv.org/pdf/1312.7651.pdf
【22】Chilimbi et.al., 10/2014 1st ver., Project Adam: Building an Efficient and Scalable Deep Learning Training System, https://www.usenix.org/system/files/conference/osdi14/osdi14-paper-chilimbi.pdf
【23】Moritz et.al., 11/2015 1st ver., SparkNet: Training Deep Networks on Spark,https://arxiv.org/pdf/1511.06051.pdf





