谷歌如何打造世界上最快的AI超级计算机系统？（下）
最近谷歌公司发表了一篇轰动人工智能系统界的论文[1]，介绍了他们如何基于MLPerf标准去刷新深度学习训练速度的世界纪录。比如，23秒完成BERT训练！28秒完成ImageNet训练！文章洋洋洒洒，从系统，算法，编译器，应用角度全方位地阐述了谷歌如何打造基于他们自己芯片的AI超级计算机以及深度学习系统。文章一共19位作者，包含了谷歌深度学习系统团队的一些专家。本文简要讨论一下谷歌公司的这篇文章，分为上下两部分。如有问题，可以通过邮箱youy@comp.nus.edu.sg联系笔者(新加坡国立大学高性能人工智能实验室主任、上海壁仞公司兼职顾问尤洋)。笔者曾在UC Berkeley读博期间在谷歌公司总部的谷歌大脑团队实习4次。
前文回顾：谷歌如何打造世界上最快的AI超级计算机系统？（上）
第3节：扩展技术
第3.1节：模型并行
对于有些模型而言，仅仅用数据并行是不够的，所以我们加上模型并行在多组TPU超级计算机上获得更高的并发性。我们用XLA的SPMD(单程序多数据)划分工具基于轻量级的标注去自动平分模型图。在图像分割模型中(例如SSD和Mask-RCNN)，我们通过标注输入图片去实现空间分割。SPMD划分工具能自动沿着空间维度去做并行计算。这些模型有相对大的空间维度(比如Mask-RCNN的空间维度是8000x1333，SSD的空间维度是300x300)。SPMD划分工具通过插入边界交换通讯操作基于空间分割计算去计算下一步所需要的响应数据。SSD和Mask-RCNN都能进行8个核上的空间划分区以实现最大化的并发度。我们又做了两点：通讯优化以及通过XLA编译器的SPMD方法(Lepikhin等人的论文[2],2020)去消除Amdahl瓶颈。这两点可以提高空间划分的并发度。例如，虽然Mask-RCNN能用最大batch size是256，但是我们却可以将Mask-RCNN的训练扩展到1024个加速器核上。
在像MLPerf的transformer基准这样的语言模型中，空间维度是很小的。所以，我们依照Shazeer等人在2018年提出的方法[3]去划分特征维度。我们实现的时候用SPMD划分工具去做标注。在这个方法里，模型的参数和响应被划分在TPU超级计算机的一个片上(超级计算机在网格视角上被划分成了很多片)。在前向传播中，每个片上的核完成局部矩阵的矩阵乘法计算。然后我们通过all-reduce操作把片上所有核上的局部响应数据整合起来用于模型下一层的计算。反向传播中有类似的局部矩阵的矩阵乘法以及之后的all-reduce去整合响应数据和梯度数据。正如模型的参数被划分那样，梯度被划分到不同核以及TPU超级计算机的每个片上，之后局部数据被整合起来。我们将在第3.3节中介绍多组超级计算机上优化梯度整合的技术。
第3.2节：参数更新划分
在传统的数据并行中，模型参数被复制并在每次训练循环结束后被优化器更新。然而，当每个核的批量数不够大的时候，计算或许会变成一个瓶颈。例如，我们以MLPerf的BERT训练为例，在512个第三代TPU芯片上，LAMB优化器的参数更新时间可以占到整个循环时间的18%。Xu等人在2020年提出了参数更新划分技术[4]，这种分布式计算技术首先执行一个reduce-scatter操作，然后使得每个加速器有整合梯度的一部分。这样每个加速器就可以算出相应的被更新的局部参数。在下一步，每个被更新的局部参数被全局广播到各个加速器，这样使得每个加速器上都有被更新的全局参数。为了获得更高的加速比，我们同时用数据并行和模型并行去处理参数更新划分。在图像分割模型中，参数是被复制的，这种情况下参数更新划分类似于数据并行。然后，当参数被分布后到不同的核之后，我们就执行多个并发的参数更新划分。
第3.3节：优化全局加法通讯操作
根据Mattson等人在2019年发表的论文[5]我们知道，在MLPerf中，梯度加法对实现高的强可扩展性效率非常重要。为了在多组TPU超级计算机上优化梯度加法通讯操作，我们在网格的Y方向上充分利用torus wrap链接。我们用一个双向的环链接在网格的Y方向上去执行reduce-scatter操作，这样的输出是沿Y方向上全局累加梯度的一部分。接下来，我们在X方向上执行reduce-scatter操作。之后我们输入局部梯度去进行参数更新。被更新的参数会先后在网格的X方向和网格的Y方向上被广播出去。值得注意的是，在数据并行中，网格X方向上的传输负载是网格Y方向上传输负载的32倍。
在MLPerf的transformer基准中，我们通过将模型参数划分到4个邻近的TPU核上去执行分布式矩阵矩阵乘法操作。这些核被放置在X方向上的一条线上。在神经网络训练的前向传播中，all-reduce操作被执行在X方向上邻居的短的环形网络链接上。然而，沿着X方向的梯度加法会跳过模型并行的邻居(因为它们划分了梯度)。图4展示了在Transformer基准中不同的环状加法操作。在BERT和Transformer模型训练中，我们也用了bfloat16浮点数运算格式去进一步减少梯度加法的开销。
图4：这个图片展示了在4个芯片上基于16x8网格的模型并行。这里展示了3个不同的环状reduction操作：(1)黑色的环状reduction用于模型并行的前向操作；(2)红色的环用于梯度块在Y方向的reduce scatter操作；(3)蓝色虚线表示在模型同类中梯度reduction操作(这里只展示了模型变好为0的情况)[1]。
第3.4节：评估指标的分布式计算
深度神经网络训练和测试在TPU加速器中执行在一个紧密的循环中。训练的结果可以更新处在TPU加速器上HBM(高速带宽内存)中的模型参数。根据MLPerf的规则，更新的参数随后每隔几个epoch就用于预测从而得到MLPerf的输出指标。在一些benchmark中，当预测的批量数大于预测数据集的样本数时，我们会在预测数据集后边加上一些仿造数据(仅仅为了便于处理数据，不影响实验结果)。在TensorFlow实现中，预测输出的张量数据被用来在TPU主机上计算预测指标(比如ResNet-50基准的Top-1精度)。然而，在JAX实现中，预测指标的计算却是全分布式的，我们通过全局加法操作得到最终的结果。
第3.5节：输入数据的流水线优化
一个扩展ResNet-50的重要挑战就是主机输入数据流水线上的负载不均衡。由于多组超级计算机的规模太大了，所以一些主机的输入数据流水线有很高的JPEG数据解压开销。我们的解决方案是把解压后的图片存在主机内存中，这样主机输入流水线只用执行3个操作：(1)随机裁剪;(2) 随机翻转；(3)按MLPerf的要求进行图片正则化使得他们有常量的平均值和方差。这会极大地增加主机输入数据流水线的吞吐率，使得它能创造一个大的预取缓冲区(prefetch buffer)。所以，主机数据流水线预处理一个大的输入图片的同时，它能向TPU输送预取缓冲区的图片。这样，我们就消除了多组TPU超级计算机输入数据流水线上的负载不均衡。
第4节：模型优化
第4.1节：BERT
用维基百科数据集去训练BERT(Devlin等人的论文[6]，2018)是MLPerf0.7版本中新加入的基准。这是一个用于语言理解的基于双向transformer结构的预训练任务。在扩展BERT的时候，我们用LAMB优化器(尤洋等人的论文[7]，2019)可以很有效地把BERT的batch size扩展到很大。这样，我们可以在4096个芯片上做有效地数据并行。这里笔者需要指出一点，笔者在ICLR2020会议上发表论文中提出LAMB优化器，当时把BERT的训练时间从3天缩短到了76分钟。本文中，谷歌也用LAMB优化器，为何训练时间只有24秒？原因有两点：一方面，谷歌确实比笔者用了更大规模的超级计算机；另一方面，更重要的是，MLPerf收敛标准跟笔者的Devlin标准是不一定的，所以MLPerf标准的计算代价本身就比Devlin标准低。
在大系统上扩展BERT需要两方面的数值优化：每步所需要的时间以及收敛需要的步数。
除了第3节提到的优化技巧之外，在模型层面，我们在减少BERT每步所需的时间上做了以下优化：减轻各体系结构单元(比如内存带宽，向量单元，以及寄存器)的压力使得TPU上的矩阵计算单元有最小的流水线开销。为了降低内存带宽的开销，我们对模型响应和梯度加法用了bfloat16浮点数运算格式，这极大地减少了每步运算所需的时间并且没有对模型的收敛产生负面影响。为了减低向量单元的压力，我们利用标量乘法和矩阵乘法的可变性把标量乘法和标量除法移到矩阵矩阵乘法维度小的那一边。为了降低寄存器溢出的频率，我们把诸如层正则化变量的一些小变量合并成一个大的TensorFlow张量。这样会大大减少储存在寄存器中的变量地址数目，因此可以减少每步训练的时间。
为了减少收敛所需的步数，我们优化了超参数以及在输入数据流水线中做了打乱数据顺序的操作。首先，我们依赖LAMB算法的扩展性用谷歌的Vizier(Golovin等人的论文[8]，2017)自动调参工具去调了大批量情况下的超参数。这样允许我们最大化数据并行，使得我们相比于模型并行用更短的时间达到相同的精度。其次，我们优化了输入数据流水线中打乱数据顺序的方法从而能保证大系统的收敛性。这一点对保证优化器的随机性很重要，因为有上千个主机的大系统往往分给每个主机的数据并不多。
第4.2节：ResNet-50
何恺明博士提出的ResNet-50模型是机器学习系统评测中最常用的模型之一。MLPerf用一千类ImageNet数据集训练ResNet-50模型作为图片分类的基准。特别地，MLPerf用了Goyal等人在2017年提出的1.5版本的ResNet-50。这个1.5版本的ResNet-50模型跟最初何恺明博士提出ResNet-50模型有轻微的差异。为了能够把ResNet-50基准扩展到多组三代TPU超级计算机上，我们应用了一系列优化技术：(1)分布式评测；(2)分布式批标准化；(3)参数更新划分；(4)梯度加法优化。MLPerf的0.7版本用了尤洋(笔者)等人提出的LARS优化器去动态扩展学习率，从而能让ResNet-50增大批量数去最大化数据并行度。在调节了动量的超参数后，我们能够用65536的批量数在88个epoch内完成训练。
第5节：性能分析
表格1展示了TensorFlow(1.x版本)和JAX在多组TPU超级计算机上完成MLPerf各个基准模型所需要的端到端训练时间。表格1也展示了我们在0.7版MLPerf的提交相对于0.6版MLPerf提交的加速比。基于这篇文章所使用的技术，我们在半分钟内完成了4个重要模型的训练。
表格1：TensorFlow(1.x版本)和JAX在多组TPU超级计算机上完成MLPerf各个基准模型所需要的端到端训练时间。这里也展示了我们在0.7版MLPerf的提交相对于0.6版MLPerf提交的加速比。
表格2比较了TensorFlow和JAX在有4096个芯片和2048个芯片的大规模分布式系统上初始化时间的差异。TensorFlow的初始化时间介于498秒到1040秒之间。然后JAX的初始化时间要低很多，只是介于122秒到294秒之间。这样因为它们在分布式编程模型的区别造成的。它们之间的权衡，性能比较，以及可用性分析详见第2节的描述。
表格2：TensorFlow和JAX在大规模分布式系统上的初始化时间对比[1]
从图5中我们可以观察到ResNet-50模型训练加速比与TPU芯片数量的关系。意料之中的是，ResNet-50取得的加速比接近最理想的扩展性。这是因为用ImageNet去训练ResNet-50的计算通讯的比值(处理一张图片的浮点运算次数与梯度中参数的个数之间的比值)非常高。
图5：ResNet-50训练加速比与芯片数目的关系[1]
为了进一步研究Amdahl定律在大系统上的瓶颈，我们在图6中展示了ImageNet/ResNet-50训练中每步所需要的计算时间以及每步所需要的all-reduce通讯时间。在图6中，蓝色区域显示的是计算时间，红色区域显示的通讯时间。这两个区域加起来就是每步训练在设备上需要的总时间。
图6：ResNet-50训练计算时间和all-reduce通讯时间与芯片数目的关系[1]
我们在图7中展示了BERT训练的加速比与TPU芯片数量的关系。从16芯片到4096芯片，BERT训练展示了极高的可扩展性。我们在图8中展示BERT训练中每步的计算时间和每步的all-reduce通讯时间。蓝色区域显示的是计算时间，红色区域显示的通讯时间。这两个区域加起来就是每步训练在设备上需要的总时间。
图7：BERT训练加速比与芯片数目的关系[1]
图8：BERT训练计算时间和all-reduce通讯时间与芯片数目的关系[1]
由于SSD，MaskR-CNN，以及Transformer三个模型的训练速度在有限的batch size的情况下依赖于模型并行。所以，我们在图9中展示了在一个主机上，模型并行产生的加速比与TPU核数之间的关系。
图9：基于TensorFlow实现，在8个TPU核上(相当于一个主机)，模型并行产生的加速比与TPU核数之间的关系[1]
当前，用于建造人工智能超级计算机的芯片主要包括第三代TPU芯片，英伟达Volta V100 GPU芯片，以及英伟达Ampere A100GPU芯片。人们分别用这三种芯片建造了人工智能超级计算机。我们在图10中比较了不同芯片建造的超级计算机对应的MLPerf基准模型训练的端到端时间。为了给大家一个公平的比较，我们在图11中展示了不同芯片建造的超级计算机对应的加速比。这里的加速比指的是相对于自身16个芯片的加速比，反应的是可扩展性。
图10：不同芯片建造的超级计算机对应的MLPerf基准模型训练的端到端时间[1]
图11：不同芯片建造的超级计算机对应的加速比。这里的加速比指的是相对于自身16个芯片的加速比，反应的是可扩展性[1]
作者简介
尤洋，壁仞科技顾问，是新加坡国立大学计算机系的助理教授(tenure-track)。他从加州大学伯克利分校计算机系获得了博士学位。他的导师是美国科学院院士，工程院院士，ACM/IEEEfellow，伯克利前计算机系主任以及EECS学院院长James Demmel教授。尤洋的研究兴趣包括高性能计算，并行算法，以及机器学习。他当前的研究重点是大规模深度学习训练算法的分布式优化。他曾创造ImageNet以及BERT训练速度的世界纪录，并被ScienceDaily，TheNextWeb，i-programmer等几十家媒体广泛报道。他设计的算法被广泛应用于谷歌，微软，英特尔，英伟达等科技巨头。尤洋近三年以第一作者身份在NIPS，ICLR，Supercomputing，IPDPS，ICS等国际重要会议或期刊上发表论文十余篇。他曾以第一作者身份获得了国际并行与分布式处理大会(IPDPS)的最佳论文(0.8%获奖率)和国际并行处理大会(ICPP)的最佳论文(0.3%获奖率)。尤洋曾获清华大学优秀毕业生，北京市优秀毕业生，国家奖学金，以及当时清华大学计算机系数额最高的西贝尔奖学金。他还在2017年获得美国计算机协会(ACM)官网上唯一颁给在读博士生的ACM-IEEECS George Michael Memorial HPC Fellowships。他也获得了颁发给伯克利优秀毕业生的LotfiA. Zadeh Prize。尤洋还被伯克利提名为ACM Doctoral Dissertation Award候选人(81名UCBerkeley EECS2020博士毕业生中选2人)。更多信息请查看他的研究小组主页(https://ai.comp.nus.edu.sg/)。
参考文献
[1]Sameer Kumar James Bradbury Cliff Young Yu Emma Wang AnselmLevskaya Blake Hechtman Dehao Chen HyoukJoong Lee Mehmet Deveci Naveen Kumar Pankaj Kanwar Shibo Wang Skye Wanderman-Milne SteveLacy Tao Wang Tayo Oguntebi Yazhou Zu Yuanzhong Xu Andy SwingEXPLORING THE LIMITS OF CONCURRENCY IN ML TRAINING ON GOOGLE TPUS[2]Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y.,Krikun,M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models withconditional computation and automatic sharding. arXiv preprintarXiv:2006.16668, 2020.[3]Shazeer, N., Cheng, Y., Parmar, N., Tran, D., Vaswani,A.,Koanantakool, P., Hawkins, P., Lee, H., Hong, M., Young, C., etal. Mesh-TensorFlow: Deep learning for supercomputers. In Advances inNeural Information Processing Systems, pp. 10414–10423, 2018.[4]Xu, Y., Lee, H., Chen, D., Choi, H., Hechtman, B., and Wang, S.Automatic cross-replica sharding of weight update in data-paralleltraining. arXiv preprint arXiv:2004.13336, 2020.[5]Mattson, P., Cheng, C., Coleman, C., Diamos, G., Micikevicius, P.,Patterson, D., Tang, H., Wei, G.-Y., Bailis, P., Bittorf, V., et al.MLPerf training benchmark. arXiv preprint arXiv:1910.01500, 2019.[6]Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert:Pre-training of deep bidirectional transformers for languageunderstanding. arXiv preprint arXiv:1810.04805, 2018.[7]You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S.,Song, X. d., Demmel, J., Keutzer, K., and Hsieh, C.-J. Large batchoptimization for deep learning: Training BERT in 76 minutes. arXivpreprint arXiv:1904.00962, 2019.[8]Golovin, D., Solnik, B., Moitra, S., Kochanski, G., Karro, J.,andSculley, D. Google Vizier: A service for black-box optimization. InProceedings of the 23rd ACM SIGKDD international conference onknowledge discovery and data mining, pp. 1487–1495, 2017.
往期推荐
壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。