上周，谷歌在论文预印本平台arxiv上发表了其关于TPUv4的深入解读论文《TPUv4:AnOpticallyReconfigurableSupercomputerforMachineLearningwithHardwareSupportforEmbeddings》(TPUv4：通过光互联可重配置的机器学习超级计算机，搭载硬件嵌入层加速)。该论文将于今年六月在ISCA2023(InternationalSymposiumonComputerArchitecture，计算机架构领域的顶级会议)上正式发表，而目前的预印本无疑为我们提供了可以一窥其全貌的机会。

为了TPU的可扩展性设计专用光学芯片，谷歌也是拼了从论文的标题可以看到，谷歌TPUv4的一个主要亮点是通过光互连实现可重配置和高可扩展性(也即标题中的“opticallyreconfigurable”)。而在论文的一开始，谷歌开门见山首先介绍的也并非传统的MAC设计、片上内存、HBM通道等AI芯片常见的参数，而是可配置的光学互联开关(reconfigurableopticalswitch)。作为论文的重中之重，这里我们也详细分析一下为什么光学互联在TPUv4设计中占了这么重要的位置，以至于谷歌甚至为了它自研了一款光学芯片。

TPUv4从一开始设计时，其目标就是极高的可扩展性，可以有数千个芯片同时加速，从而实现一个为了机器学习模型训练而设计的超级计算机。在谷歌的设计中，超级计算机的拓扑结构为：将4x4x4(64)个TPUv4芯片互联在一起形成一个立方体结构(cube)，然后再把4x4x4这样的cube用连在一起形成一个总共有4096个TPUv4的超级计算机。

在这样的拓扑中，物理距离较近的TPUv4(即在同一个4x4x4cube中的芯片)可以用常规的电互联(例如铜绞线)方法连接，但是距离较远的TPU之间(例如在cube之间的互联)就必须使用光互连，原因就在于在如此大规模的超级计算机中，芯片之间的数据互联在很大程度上会决定整体计算的效率;如果数据互联效率不够高的话，很多时候芯片都在等待来自其他芯片数据到达以开始计算，这样就形成了效率浪费。为了避免这样“芯片等数据”的情形出现，就必须确保芯片之间互联能拥有高带宽，低延迟。而光互连对于物理距离较远的芯片就成为了首选。

光互连在高性能计算中的使用也并非新闻，而谷歌在TPUv4中的主要突破是使用可重配置的光互连(即加入光路开关，opticalcircuitswitchOCS)来快速实现不同的芯片互联拓扑。换句话说，芯片之间的互联并非一成不变的，而是可以现场可重配置的。这样做可以带来许多好处，其中最主要的就是可以根据具体机器学习模型来改变拓扑，以及改善超级计算机的可靠性。

从拓扑结构来说，不同的机器学习模型对于数据流的要求大致可以分为三类，即数据平行(每块芯片都加载整个模型，不同的芯片处理数据集中不同的数据)，模型并行(模型中有些层特别大，因此每块芯片只负责这样很大的层中的一部分计算)，以及流水线并行(把模型中的不同层交给不同的芯片计算)，而不同的数据流就对应了不同的TPU互联拓扑。当有了可重配置光互连之后，就可以根据具体模型数据流来调整TPU之间的互联拓扑，从而实现最优的性能，其提升可超过2倍。

另一个优势就是可靠性。在这样拥有海量芯片组成的超级计算机中，一个重要的考量就是，如果有一小部分的芯片不工作了，如何确保整体超级计算机仍然能维持较高的性能?如果使用常规的固定互联架构，那么一个芯片出故障可能会影响整个系统工作。而在有了可重配置的光互连之后，需要做的只需要把出故障的芯片绕过，就不会影响整个系统的工作，最多会牺牲一点整体的性能。谷歌在论文中给出了一个单芯片故障率和系统平均性能影响的曲线图，在使用可配置光互连(以及光路开关)时，假设芯片可靠率在99%的情况下，其整体系统的平均性能提升比不使OCS可高达6倍，可见光互连开关的重要性。

为了实现数据中心级的可配置光互连，需要光路开关首先能高效扩展到超高数量的互联数(例如1000x1000)，同时需要实现低开关切换延迟，低成本，以及低信号损耗。谷歌认为现有的商用方案都不够满意，因此谷歌的做法是自研了一款光路开关芯片Palomar，并且使用该芯片实现了全球首个数据中心级的可配置光互连，而TPUv4就是搭配了这款自研光路开关芯片从架构上实现了高性能。谷歌自研的光路开关芯片Palomar使用的是基于MEMS反射镜阵列的技术，具体原理是使用一个2DMEMS反射镜阵列，通过控制反射镜的位置来调整光路，从而实现光路的切换。使用MEMS的光路开关芯片可以实现低损耗，低切换延迟(毫秒级别)以及低功耗。在经过一系列优化之后，光路系统的成本也控制得很低，在整个TPUv4超级计算机成本中占5%以下。

算法-芯片协同设计是TPUv4的灵魂如果说可重配置光互联给TPUv4提供了良好的根基的话，那么算法-芯片协同设计就是TPUv4的灵魂。算法-芯片协同设计包括两部分，一部分是如何根据算法优化芯片，而另一方面是如何根据芯片去优化算法，在TPUv4的架构中，两者都得到了仔细考虑。

我们首先分析TPUv4如何根据算法来优化芯片。如果说2017年TPUv1发表时候，其主要解决的还是卷积神经网络CNN的话，那么在2023年来看，CNN的加速问题早已经被既觉得差不多了，更多的是如何处