上周，谷歌在论文预印本平台arxiv上发表了其关于TPU v4的深入解读论文《TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings》（TPU v4：通过光互联可重配置的机器学习超级计算机，搭载硬件嵌入层加速）。
该论文将于今年六月在ISCA 2023（International Symposium on Computer Architecture，计算机架构领域的顶级会议）上正式发表，而目前的预印本无疑为我们提供了可以一窥其全貌的机会。
为了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超级计算机（由4096个TPU v4组成）拓扑结构，图/谷歌
在这样的拓扑中，物理距离较近的TPU v4（即在同一个4x4x4 cube中的芯片）可以用常规的电互联（例如铜绞线）方法连接，但是距离较远的TPU之间（例如在cube之间的互联）就必须使用光互连，原因就在于在如此大规模的超级计算机中，芯片之间的数据互联在很大程度上会决定整体计算的效率；
如果数据互联效率不够高的话，很多时候芯片都在等待来自其他芯片的数据到达以后开始计算，这样就形成了效率浪费。为了避免这样“芯片等数据”的情形出现，就必须确保芯片之间互联能拥有高带宽，低延迟。而光互连对于物理距离较远的芯片就成为了首选。
光互连在高性能计算中的使用也并非新闻，而谷歌在TPU v4中的主要突破是使用可重配置的光互连（即加入光路开关，optical circuit switch OCS）来快速实现不同的芯片互联拓扑。
换句话说，芯片之间的互联并非一成不变的，而是可以现场可重配置的。这样做可以带来许多好处，其中最主要的就是可以根据具体机器学习模型来改变拓扑，以及改善超级计算机的可靠性。