最近谷歌公司发表了一篇轰动人工智能系统界的论文[1]，介绍了他们如何基于MLPerf标准去刷新深度学习训练速度的世界纪录。比如，23秒完成BERT训练！28秒完成ImageNet训练！文章洋洋洒洒，从系统，算法，编译器，应用角度全方位地阐述了谷歌如何打造基于他们自己芯片的AI超级计算机以及深度学习系统。文章一共19位作者，包含了谷歌深度学习系统团队的一些专家。本文简要讨论一下谷歌公司的这篇文章，分为上下两部分。如有问题，可以通过邮箱youy@comp.nus.edu.sg联系笔者(新加坡国立大学高性能人工智能实验室主任、壁仞科技顾问尤洋)。笔者曾在UC Berkeley读博期间在谷歌公司总部的谷歌大脑团队实习4次。

第1节：介绍

以自然语言处理为代表的深度学习革命正处在类似于美苏60年代“太空竞赛”的境地，导致模型的大小正以爆炸性的速度增长。显然，大模型带来了极大的好处，比如OpenAI的超大模型GPT-3已经可以在轻微的人类引导下去生成一些有可读性的散文，效果明显是之前的模型无法达到的。同时OpenAI研究院也观察到，世界最佳模型所消耗的计算资源每3.5个月就要翻一番。这比当时的摩尔定律夸张多了。之前摩尔定律表明每18月单位面积的晶体管数量翻一番。由摩尔定律导致的安迪-比尔定律直接地预示了计算机产业几十年的辉煌。

所以，为了满足这些大模型的快速训练需求，我们需要更大的超级计算机。这场深度学习革命引发的超算革命给GPU带来了巨额的投资：过去几年，英伟达公司的营收见证过50%的年增长率。英伟达公司的股票也从2012年11月的11美元增长到了2020年11月582美元(53倍)。

2015年，谷歌的TPU神经网络专用加速器横空出世。谷歌宣称，每个TPU芯片比之前的神经网络芯片在功耗比，低延迟效率，以及最高性能上提升了10倍(Jouppi等人的论文[2],2017)。在之后的两年，谷歌基于二代TPU的256芯片超计算机就展示出了完美的并行效率(Jouppi等人的论文[3],2020)。之后谷歌基于3代TPU将超级计算机的规模做到了1024芯片的规模(Kumar等人的论文[4],2019)。英伟达和其它GPU供应商也在建造类似规模的AI超级计算机。微软和OpenAI甚至要打造一个一万个GPU的超级计算机(Langston的文章[5],2020)。这场类似于“美苏星球大战”的比拼将会用精准的模型去试图逼近通用人工智能。这一点是否能实现还有待观察。但是，无疑的是，这场比拼中的各个巨头对硬件投入是雄心勃勃的。

美苏太空竞赛可以用类似于“绕地飞行”或“登陆月球”等一些里程碑事件去标记成就。然而，这次AI超算竞争却在开始的时候没有具体衡量指标。所以，谷歌、英伟达、英特尔、阿里巴巴、华为等工业界巨头意识到这个问题后提出了MLPerf标准(MLPerf.org)。特别地，MLPerf训练标准吸引了众多有高性能计算能力的实体去用超级计算机以最短的时间完成神经网络的训练(Mattson等人的论文[6],2019)。

事实证明，在过去的两年，MLPerf标准对深度学习社区产生了积极而深远的影响。因为刷新MLPerf标准的努力和创新产生了新的系统优化技术，更实用的代码库，更高效的编译器，以及最合适的应用层代码。

常规的TPU超级计算机只有1024个芯片。我们在上次的MLPerfv0.6版本比赛中只用了常规的TPU超级计算机。为了探索MLPerf模型并行计算的极限，我们用第三代TPU芯片组装了一个4096芯片的多组超级计算机(详见图1)。两个TPU超级计算机在网格结构的X方向上被光纤连接起来(详见图2)。这些组间的光纤链接要比普通的组内光纤链接长。我们用了128x32的二维网状结构将所有的芯片连接起来用于训练MLPerf的模型。因为第三代TPU芯片的路由表只能容纳1024个地址，所以我们用了一种稀疏的路由策略。在这种策略下，每个芯片只能看到与它同一行或同一列的芯片。这个策略足以使all-reduce操作达到峰值的通讯传输效率。

图1：常规TPU超级计算与多组TPU超级计算机[1]

图2：4组TPU超级计算机的配置，跨组链接在横向连接了两个TPU超级计算机[1]

第2节：多种编程框架

尽管TPU的主要前端编程框架一直是TensorFlow(Abadi等人的论文[7],2016)，TPU硬件和XLA编译器其实是通用的工具：它们是可以支持其它的编程框架。因此，在这篇论文中，我们选择给TensorFlow和JAX都做评测。JAX(Frostig等人的论文[8],2018)是一种面向研究的基于XLA的数值计算编程框架。这两种编程模型都需要额外的软件工程在多组超级计算机上获得高效的扩展性。但是它们最终可以获得相似的评测结果。

如图3所示，由于TensorFlow和JAX在架构上的不同，它们在大规模运算上也有性能差异。首先，它们有不同分阶段方法。TensorFlow在Python中嵌入了一种表达式的动态即时编程语言，然后用XLA去即时编译图的子集(TensorFlow的图可以分布在TPU加速器上和CPU上)。相比之下，JAX少了一个阶段：它是一个在Python中嵌入了即时编译的XLA程序的编程框架。其中，即时编译的XLA程序用于在加速器上的静态编译性能以及加速器网络上的并行性。JAX用于动态以及加速计算。所以，TensorFlow比JAX多了一个编译阶段，我们用多线程去加速了这个多出来的编译阶段。同时，JAX也需要我们更加小心地去处理Python的瓶颈。比如，把一些类似于数据读取的干扰性任务移出主线程。

图3：TensorFlow和JAX编程模型在第三代TPU上使用的示意图[1]

其次，TensorFlow和JAX可以支持不同的分布式编程模型。JAX采用了一种多客户端的方式去做分布式编程，使得超级计算机中每个主机分别有一份同样的JAX代码(包括Python解释器)。JAX程序只在两个地方进行通讯：(1)在初始阶段去设置TPU的网格连接。(2)在模型训练过程中在网络上进行XLA编译好的all-reduce操作。与之相反的是，TensorFlow采用单客户端的方式在TPU上进行编程。这种方式会给一个Python进程全局视角并允许它控制整个分布式系统。这个Python进程运行在超级计算机的一个主机上。这个主机将TensorFlow的图均分后通过RPC在网络上发给剩下的主机执行。

TensorFlow和JAX在实用性和性能特点上也有不同之处。尽管TensorFlow的单客户端分布式系统能够让用户代码直接控制整体的运行负载，JAX的多客户端方法使得代码可以直接控制各个独立的计算单元。JAX在各个主机上独立地调用XLA编译器，这依赖于确定性编译技术去避免不同主机程序的不兼容性。然而，TensorFlow只编译一次并且把二进制文件分给所有主机节点。TensorFlow的这种多设备图其实能引起Amdahl定律瓶颈。原因是因为客服端进程在图构造和优化时间上的开销是跟节点数成正比的。然后，除了TPU拓扑网格链接的初始化之外，JAX的设置时间是不随着节点数增加而增加的。所以，JAX在这一点上看似有更好的可扩展性。

作者简介

尤洋，壁仞科技顾问，是新加坡国立大学计算机系的助理教授(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名UC Berkeley EECS2020博士毕业生中选2人)。更多信息请查看他的研究小组主页(https://ai.comp.nus.edu.sg/)。

敬请期待：谷歌如何打造世界上最快的AI超级计算机系统？（下）
 

参考文献

[1] Sameer Kumar James Bradbury Cliff Young Yu Emma Wang Anselm Levskaya Blake Hechtman  Dehao Chen  HyoukJoong Lee Mehmet Deveci Naveen Kumar Pankaj Kanwar  Shibo Wang Skye Wanderman-Milne Steve Lacy Tao Wang  Tayo Oguntebi  Yazhou Zu  Yuanzhong Xu Andy Swing EXPLORING THE LIMITS OF CONCURRENCY IN ML TRAINING ON GOOGLE TPUS

[2] Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., et al. In-datacenter performance analysis of a tensor processing unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture, pp. 1–12, 2017.

[3] Jouppi, N. P., Yoon, D. H., Kurian, G., Li, S., Patil, N., Laudon, J., Young, C., and Patterson, D. A domainspecific supercomputer for training deep neural networks. Communications of the ACM, 63(7):67–78, 2020.

[4] Kumar, S., Bitorff, V., Chen, D., Chou, C., Hechtman, B., Lee, H., Kumar, N., Mattson, P., Wang, S., Wang, T., et al. Scale MLPerf-0.6 models on Google TPU-v3 pods. arXiv preprint arXiv:1909.09756, 2019.

[5] Langston, J. Microsoft announces new supercomputer, lays out vision for future AI work. MicroSoft Blog, 2020.

[6] 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.

[7] Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., et al. TensorFlow: A system for large-scale machine learning. In 12th {USENIX} symposium on operating systems design and implementation ({OSDI} 16), pp. 265–283, 2016.

[8] Frostig, R., Johnson, M. J., and Leary, C. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 2018.