谷歌TPUv7挑战AI芯片王座（六）：揭秘ICI超大规模互联与软件战略的破局之道

为何采用谷歌的ICI 3D环面架构？

除了可以让人花费无数小时绘制的精美立方体示意图之外，谷歌独特的ICI纵向扩展网络究竟有何卓越之处？

系统规模： 最显而易见的好处是TPUv7 Ironwood所支持的、高达9,216个TPU的极大系统规模。尽管由于有效吞吐率下降的缺点，9,216的最大切片规模可能极少被使用，但包含数千个TPU的切片能够且确实被普遍采用。这远大于商业加速器市场及其他定制芯片供应商常见的64或72个加速器的系统规模。

可重构性与可互换性： 使用光路交换器意味着网络拓扑天生支持重新配置网络连接，以支持大量不同的拓扑结构——理论上可达数千种。谷歌的文档网站列出了10种不同的组合（本节前文图片），但这些仅是最常见的3D切片形状——实际可用的远不止于此。

即使是相同规模的切片也能以不同方式重新配置。在下图所示的简单“扭转2D环面”示例中，我们可以看到，通过循环连接到不同X坐标索引而非相同X坐标索引，可以减少最坏情况下的跳数和最坏情况下的二分带宽。这有助于提升“全体到全体”集合操作的吞吐量。一个TPUv7集群将在4x4x4立方体级别进行扭转。

来源：SemiAnalysis，谷歌

可重构性也为实现多种并行模式打开了大门。在64或72个加速器的系统规模中，不同的并行组合通常受限于64的因数。而对于ICI纵向扩展网络，为实现精确匹配所需的数据并行、张量并行和流水线并行组合而实施拓扑的可能性则非常丰富。

光路交换器允许将任何立方体的任何“+”面连接到任何其他立方体的“-”面，这一事实意味着立方体具备完全的可互换性。切片可以由任意一组立方体构成。因此，若出现任何故障或用户需求、使用情况发生变化，这都不会妨碍新拓扑切片的形成。

来源：谷歌

更低成本： 谷歌的ICI网络比大多数基于交换机的纵向扩展网络成本更低。尽管所使用的全双工光模块因采用了环形器而可能略显昂贵，但网状网络减少了对交换机和端口的总数量需求，并消除了交换机间连接所产生的成本。

来源：SemiAnalysis

低延迟与更优的局部性： 在TPU间使用直接链路意味着，对于物理位置彼此靠近或经重新配置后直接相互连接的TPU，有可能实现更低的延迟。彼此靠近的TPU也具备更好的数据局部性。

数据中心网络——超越9216个TPU的扩展

数据中心网络是一个独立于ICI的网络，其作用兼具典型的后端和前端网络。它在更大的范围内进行连接——就TPUv7集群而言，可达14.7万个TPU。

正如我们早先关于“阿波罗计划”的帖子所讨论的，谷歌提出用Paloma光路交换器取代传统“Clos”架构中包含电子分组交换机的骨干层。谷歌的DCN由一个光交换数据中心网络互连层构成，该层结合了多个聚合区块，每个聚合区块连接多个9216个TPU的ICI集群。

2022年，谷歌的“阿波罗”项目提出了一种DCN架构，描述了对TPUv4 Pod（每个Pod规模为4096个TPU）使用136x136的光路交换器。DCNI层的光路交换器被组织成4个阿波罗区域，每个区域最多包含8个机架，每个机架有8台光路交换器，总计256台光路交换器。对于Ironwood，为了在同一网络上支持多达14.7万个TPUv7，我们推测光路交换器的端口数将接近翻倍，而非增加光路交换器的最大数量。

下图展示了一个使用32个机架、容纳256台300x300光路交换器的Ironwood DCN网络可能的样子。假设每个聚合区块的骨干层之间不存在超额订阅，DCN最多可连接16个ICI Pod，由4个聚合区块各连接4个ICI Pod——总计147456个TPU。DCNI层连接这4个聚合区块——如下图所示的最顶层。与ICI类似，采用全双工光模块连接至光路交换器，以最大化每台光路交换器上每个端口的带宽。

虽然现有的Ironwood集群可能只有1个或2个聚合区块，但谷歌DCN独特的架构允许将新的TPU聚合区块添加到网络中，而无需进行大规模的重新布线。

通过在DCNI层使用光路交换器，DCN结构的规模可以逐步扩展，并且网络可以重新分配资源以支持新的聚合区块。此外，聚合区块的带宽可以升级，而无需改变DCN层的构成。这使得现有聚合区块的链路速度可以在不改变网络基本架构的情况下得到更新。当然，这种结构扩展过程不能无限进行——在规模极大时，重新配置网络将变得难以管理。

来源：SemiAnalysis，谷歌

TPU软件战略——另一场重大转变

传统上，TPU软件和硬件团队主要面向内部服务。这带来了一些优势，例如没有市场营销团队为夸大理论浮点算力而施加的压力。

仅面向内部的另一个优势是，TPU团队高度优先处理内部功能请求和优化内部工作负载。其劣势在于，他们不太关心外部客户或外部工作负载。TPU生态系统中的外部开发者数量远低于CUDA生态系统。这是TPU以及所有非英伟达加速器的主要弱点之一。

此后，谷歌已修订其面向外部客户的软件战略，并对TPU团队的关键绩效指标及其对AI/ML生态系统的贡献方式进行了重大调整。我们将讨论两项主要变革：

1.为PyTorch提供“原生”TPU支持的大规模工程投入

2.为vLLM/SGLang提供TPU支持的大规模工程投入

通过观察谷歌在各个TPU软件代码库的贡献数量，外部化战略清晰可见。我们可以看到，从三月份开始，对vLLM的贡献显著增加。随后从五月份起，创建了“tpu-inference”代码库，这是官方的vLLM TPU统一后端，自此相关活动十分活跃。

来源：GitHub, SemiAnalysis

传统上，谷歌仅对Jax/XLA:TPU堆栈提供一流支持（以及TensorFlow/TF-Mesh），但将TPU上的PyTorch视为二等公民。它依赖于通过PyTorch/XLA进行惰性张量图捕获，而非提供一流的热切执行模式。此外，它不支持PyTorch原生分布式API，也不支持PyTorch原生并行API，而是依赖于奇怪的、非核心的XLA SPMD API。这导致习惯了GPU上原生PyTorch CUDA后端并试图转向TPU的外部用户体验不佳，感觉不够原生。

来源：XLA

原文媒体：SemiAnalysis

原文作者：Dylan Patel; Myron Xie; Daniel Nishball

原文链接：

https://newsletter.semianalysis.com/p/tpuv7-google-takes-a-swing-at-the