相比较使用英伟达 A100 所创建的超级计算机,使用谷歌 TPUv4 构建的超级计算机速度快 1.2-1.7 倍,功耗降低 1.3-1.9 倍。
目前在人工智能计算领域,英伟达占据主导地位,超过 90% 的人工智能项目开发都使用其芯片。而谷歌试图在该市场上和英伟达竞争,TPU 的全称为 Tensor Processing Units,是谷歌自研的芯片。
为了TPU的可扩展性设计专用光学芯片,谷歌也是拼了
从论文的标题可以看到,谷歌TPU v4的一个主要亮点是通过光互连实现可重配置和高可扩展性(也即标题中的“optically reconfigurable”)。而在论文的一开始,谷歌开门见山首先介绍的也并非传统的MAC设计、片上内存、HBM通道等AI芯片常见的参数,而是可配置的光学互联开关(reconfigurable optical switch)。作为论文的重中之重,这里我们也详细分析一下为什么光学互联在TPU v4设计中占了这么重要的位置,以至于谷歌甚至为了它自研了一款光学芯片。
TPU v4从一开始设计时,其目标就是极高的可扩展性,可以有数千个芯片同时加速,从而实现一个为了机器学习模型训练而设计的超级计算机。在谷歌的设计中,超级计算机的拓扑结构为:将4x4x4(64)个TPU v4芯片互联在一起形成一个立方体结构(cube),然后再把4x4x4这样的cube用连在一起形成一个总共有4096个TPU v4的超级计算机。
在这样的拓扑中,物理距离较近的TPU v4(即在同一个4x4x4 cube中的芯片)可以用常规的电互联(例如铜绞线)方法连接,但是距离较远的TPU之间(例如在cube之间的互联)就必须使用光互连,原因就在于在如此大规模的超级计算机中,芯片之间的数据互联在很大程度上会决定整体计算的效率;如果数据互联效率不够高的话,很多时候芯片都在等待来自其他芯片数据到达以开始计算,这样就形成了效率浪费。为了避免这样“芯片等数据”的情形出现,就必须确保芯片之间互联能拥有高带宽,低延迟。而光互连对于物理距离较远的芯片就成为了首选。
光互连在高性能计算中的使用也并非新闻,而谷歌在TPU v4中的主要突破是使用可重配置的光互连(即加入光路开关,optical circuit switch OCS)来快速实现不同的芯片互联拓扑。换句话说,芯片之间的互联并非一成不变的,而是可以现场可重配置的。这样做可以带来许多好处,其中最主要的就是可以根据具体机器学习模型来改变拓扑,以及改善超级计算机的可靠性。
从拓扑结构来说,不同的机器学习模型对于数据流的要求大致可以分为三类,即数据平行(每块芯片都加载整个模型,不同的芯片处理数据集中不同的数据),模型并行(模型中有些层特别大,因此每块芯片只负责这样很大的层中的一部分计算),以及流水线并行(把模型中的不同层交给不同的芯片计算),而不同的数据流就对应了不同的TPU互联拓扑。当有了可重配置光互连之后,就可以根据具体模型数据流来调整TPU之间的互联拓扑,从而实现最优的性能,其提升可超过2倍。
另一个优势就是可靠性。在这样拥有海量芯片组成的超级计算机中,一个重要的考量就是,如果有一小部分的芯片不工作了,如何确保整体超级计算机仍然能维持较高的性能?如果使用常规的固定互联架构,那么一个芯片出故障可能会影响整个系统工作。而在有了可重配置的光互连之后,需要做的只需要把出故障的芯片绕过,就不会影响整个系统的工作,最多会牺牲一点整体的性能。谷歌在论文中给出了一个单芯片故障率和系统平均性能影响的曲线图,在使用可配置光互连(以及光路开关)时,假设芯片可靠率在99%的情况下,其整体系统的平均性能提升比不使OCS可高达6倍,可见光互连开关的重要性。
为了实现数据中心级的可配置光互连,需要光路开关首先能高效扩展到超高数量的互联数(例如1000x1000),同时需要实现低开关切换延迟,低成本,以及低信号损耗。谷歌认为现有的商用方案都不够满意,因此谷歌的做法是自研了一款光路开关芯片Palomar,并且使用该芯片实现了全球首个数据中心级的可配置光互连,而TPU v4就是搭配了这款自研光路开关芯片从架构上实现了高性能。谷歌自研的光路开关芯片Palomar使用的是基于MEMS反射镜阵列的技术,具体原理是使用一个2D MEMS反射镜阵列,通过控制反射镜的位置来调整光路,从而实现光路的切换。使用MEMS的光路开关芯片可以实现低损耗,低切换延迟(毫秒级别)以及低功耗。在经过一系列优化之后,光路系统的成本也控制得很低,在整个TPU v4超级计算机成本中占5%以下。
Google自研TPU,五年更新四代
Google最早于2016年宣布首款内部定制的AI芯片,区别于训练和部署AI模型的最常见的组合架构,即CPU和GPU组合,第一代TPU在那场世界著名的人机围棋大战助力AlphaGo打败李世石“一战成名”,宣告并不是只有GPU才能做训练和推理。
Google第一代TPU采用28nm工艺制程,功耗大约40W,仅适用于深度学习推理,除了AlphaGo,也用在Google搜索、翻译等机器学习模型中。
2017年5月,Google发布了能够实现机器学习模型训练和推理的TPU v2,达到180TFLOPs浮点运算能力,同时内存带宽也得以提升,比同期推出的CPU AI工作负载提升30倍,比GPU AI工作负载提升15倍,被基于4块TPU v2的AlphaGo击败的世界围棋冠军柯洁最直观地感受了这一切。
2018年5月,Google又发布第三代TPU v3,性能是上一代TPU的两倍,实现420TFLOPs浮点运算,以及128GB的高带宽内存。
按照一年一次迭代更新的节奏,Google理应在2019年推出第四代TPU,不过这一年的I/O大会上,Google推出的是第二代和第三代TPU Pod,可以配置超过1000颗TPU,大大缩短了在进行复杂的模型训练时所需耗费的时间。
在AI芯片发展史上,无论是从片上内存上,还是从可编程能力来看,Google TPU都是不可多得的技术创新,打破GPU的“垄断”地位,且打开云端AI芯片的新竞争格局。
发展五年的Google TPU在今天依然保持着强劲的竞争力,未来的世界是什么样的?Google TPU已经告诉了我们一小部分答案。
TPU v4与人工智能芯片的未来
从TPU v4的设计中,我们可以看到人工智能芯片未来的一些方向,而这些方向是我们在Nvidia的GPU等其他主流人工智能芯片的设计中也看到的:
首先就是对于高效互联和规模化的支持。随着人工智能模型越来越大,对于这类模型的支持主要依赖人工智能芯片的可扩展性(即如何让多芯片可以高效并可靠地一起分工合作来加速这样的大模型),而不是一味提高单芯片的能力来支持大模型,因为模型的演进总是要比芯片的设计迭代更快。在这个领域,不同的芯片公司会有不同的侧重,例如AMD侧重较为微观层面的使用chiplet来实现封装级别的可扩展性,Nvidia有NvLink等芯片技术来实现单机多卡之间的可扩展性和性能提升,谷歌则直接为了海量TPU互联设计了一款光路开关芯片;但是这些公司之间的共性,即对于人工智能芯片可扩展性的支持以满足大模型的需求,却是相当一致的。从这个角度来看,未来可扩展性(例如数据互联带宽)有可能会成为与峰值算力一样的人工智能芯片主要指标,而这也让人工智能芯片设计更加跨界:即不仅仅是需要对于数字逻辑和计算机架构方面的资源,同时也需要在封装、数据互联等领域都有积累。
此外,人工智能芯片与算法之间的结合继续保持紧密关系,算法-芯片协同设计仍然将是未来人工智能芯片继续提升性能的主要手段之一。我们目前已经看到了谷歌、Nvidia等在算法-芯片协同设计中的大量成果:包括对于新的数制(Nvidia的FP16、FP8,谷歌的BF16等)的支持,对于计算特性的支持(Nvidia对于稀疏计算的支持),以及对于模型关键算法的直接专用加速器的部署(Nvidia的transformer acclerator,谷歌的SC等)。随着摩尔定律未来越来越接近物理极限,预计未来人工智能芯片性能进一步提升会越来越倚赖算法-芯片协同设计,而另一方面,由于有算法-芯片协同设计,我们预计未来人工智能芯片的性能仍然将保持类似摩尔定律的接近指数级提升,因此人工智能芯片仍然将会是半导体行业未来几年最为热门的方向之一,也将会成为半导体行业未来继续发展的重要引擎。