谷歌称自家AI超算系统优于英伟达，秘密就在最新TPU v4芯片中

当地时间周二（4月4日），Alphabet旗下的谷歌公司公布了其用于训练人工智能（AI）模型的超级计算机的最新细节。该公司称，这些系统比英伟达公司的同期系统更快、更节能。

据了解，谷歌张量处理器（tensorprocessingunit，简称TPU芯片）是该公司为机器学习定制的专用芯片（ASIC），第一代发布于2016年，为当时大热的AlphaGo提供算力。与GPU相比，TPU采用低精度计算，在几乎不影响深度学习处理效果的前提下大幅降低了功耗、加快运算速度。同时，TPU使用了脉动阵列等设计来优化矩阵乘法与卷积运算。

GoogleCEOSundarPichai在GoogleI/O2021上宣布TPUv4。（来源：Google）

作为图灵奖得主、计算机架构巨擘，大卫・帕特森（DavidPatterson）在2016年从UCBerkeley退休后，以杰出工程师的身份加入了谷歌大脑团队，为几代TPU的研发做出了卓越贡献。

谷歌表示，目前公司90%以上的AI训练工作都是通过谷歌自主研制的TPU芯片完成的。其AI训练的过程是：通过给模型输入数据，使其能够用类似人类的文本回应、并具备生成图像等功能。

当前谷歌使用的是最新的第四代TPU芯片（TPUv4），谷歌公司周二发布的一篇题为《TPUv4:AnOpticallyReconfigurableSupercomputerforMachineLearningwithHardwareSupportforEmbeddings》的论文，详细阐述了该公司是如何利用自己定制的光路交换器（OCS），将4000多个芯片组合到一台超级计算机中，从而帮助连接一台台独立机器。谷歌表示，其已经部署了数十台采用第四代TPU打造的超级计算机，供内部使用和外部通过谷歌云使用。

相关论文

改善芯片间的连接成AI超算关键竞争点

随着OpenAI的ChatGPT、谷歌的Bard等AI聊天机器人的竞争愈演愈烈，改善芯片间的连接已成为开发AI超算的科技公司的一个关键竞争点。

由于ChatGPT、Bard等产品背后使用的大型语言模型的规模太庞大了，动辄千亿的参数量无法存储在单一芯片上。所以，这些模型必须分散在数千个芯片上，然后这些芯片必须同时工作，花费数周或更长时间来训练语言模型。

谷歌的PaLM是该公司迄今为止公开披露的最大的语言模型。其训练方式是：通过将其分配到两台各包含4000个芯片的超级计算机中，历时50天进行训练。

谷歌周二称，其通过光电路交换机，其超级计算机可以轻松地动态重新配置芯片之间的连接，有助于避免出现问题并实时调整以提高性能。

下图展示了TPUv44×3方式6个「面」的链接。每个面有16条链路，每个块总共有96条光链路连接到OCS上。要提供3D环面的环绕链接，相对侧的链接必须连接到相同的OCS。因此，每个4×3块TPU连接到6×16÷2=48个OCS上。PalomarOCS为136×136（128个端口加上8个用于链路测试和修复的备用端口），因此48个OCS连接来自64个4×3块（每个64个芯片）的48对电缆，总共并联4096个TPUv4芯片。

(来源：Google)

根据这样的排布，TPUv4（中间的ASIC加上4个HBM堆栈）和带有4个液冷封装的印刷电路板(PCB)。该板的前面板有4个顶部PCIe连接器和16个底部OSFP连接器，用于托盘间ICI链接。

(来源：Google)

随后，八个64芯片机架构成一台4096芯片超算。

与超级计算机一样，工作负载由不同规模的算力承担，称为切片：64芯片、128芯片、256芯片等。下图显示了当主机可用性从99.0%到99.9%不等有，及没有OCS时切片大小的「有效输出」。如果没有OCS，主机可用性必须达到99.9%才能提供合理的切片吞吐量。对于大多数切片大小，OCS也有99.0%和99.5%的良好输出。

(来源：Google)

与Infiniband相比，OCS的成本更低、功耗更低、速度更快，成本不到系统成本的5%，功率不到系统功率的3%。每个TPUv4都包含SparseCores数据流处理器，可将依赖嵌入的模型加速5至7倍，但仅使用5%的裸片面积和功耗。

该公司的研究员NormJouppi和工程师DavidPatterson在上述论文中写道，“电路切换可以很容易地绕过故障组件…这种灵活性甚至允许我们改变超级计算机互连的拓扑结构，以加速机器学习模型的性能。”

稀疏核（SparseCore，SC）设计

在新论文上，谷歌着重介绍了稀疏核（SparseCore，SC）的设计。在大模型的训练阶段，嵌入训练可以放在TensorCore或超级计算机的主机CPU上处理。TensorCore具有宽VPU和矩阵单元，并针对密集操作进行了优化。由于小的聚集/分散内存访问和可变长度数据交换，在TensorCore上放置嵌入其实并不是最佳选择。在超级计算机的主机CPU上放置嵌入会在CPUDRAM接口上引发阿姆达尔定律瓶颈，并通过4:1TPUv4与CPU主机比率放大。数据中心网络的尾部延迟和带宽限制将进一步限制训练系统。

对此，谷歌认为可以使用TPU超算的总HBM容量优化性能，加入专用ICI网络，并提供快速收集/分散内存访问支持。这导致了SparseCore的协同设计。

SC是一种用于嵌入训练的特定领域架构(DSA)，从TPUv2开始，后来在TPUv3和TPUv4中得到改进。SC相对划算，只使用芯片面积和功率的约5%左右。SC结合超算规模的HBM和ICI来创建一个平坦的、全局可寻址的内存空间（TPUv4中为128TiB）。与密集训练中大参数张量的全部归约相比，较小嵌入向量的全部传输使用HBM和ICI以及更细粒度的分散/聚集访问模式。

作为独立的核心，SC允许跨密集计算、SC和ICI通信进行并行化。下图显示了SC框图，谷歌将其视为「数据流」架构（dataflow），因为数据从内存流向各种直接连接的专用计算单元。

(来源：Google)

最通用的SC单元是16个计算块（深蓝色框）。每个tile都有一个关联的HBM通道，并支持多个未完成的内存访问。每个tile都有一个FetchUnit、一个可编程的8-wideSIMDVectorProcessingUnit和一个FlushUnit。获取单元将HBM中的激活和参数读取到2.5MiB稀疏向量内存(Spmem)的图块切片中。scVPU使用与TC的VPU相同的ALU。FlushUnit在向后传递期间将更新的参数写入HBM。此外，五个跨通道单元（金色框）执行特定的嵌入操作，正如它们的名称所解释的那样。

与TPUv1一样，这些单元执行类似CISC的指令并对可变长度输入进行操作，其中每条指令的运行时间都取决于数据。

(来源：Google)

在特定芯片数量下，TPUv3/v4对分带宽比高2-4倍，嵌入速度可以提高1.1-2.0倍。

下图展示了谷歌自用的推荐模型（DLRM0）在不同芯片上的效率。TPUv3比CPU快9.8倍。TPUv4比TPUv3高3.1倍，比CPU高30.1倍。谷歌探索了TPUv4超算用于GPT-3大语言模型时的性能，展示了预训练阶段专家设计的1.2倍改进。

(来源：Google)

新版本TPU与英伟达A100对比

虽然谷歌直到现在才披露关于其超级计算机的详细信息，不过事实上，自2020年以来，该系统已经在公司内部上线，应用于该公司位于俄克拉荷马州的一个数据中心。

来自位于俄克拉荷马州的谷歌ML集群的TPUv4pod的八分之一，该公司声称该集群使用约90%的无碳能源运行。（来源：Google）

据悉，初创公司Midjourney使用TPUv4系统来训练其模型，在输入几句文字后，模型就会生成新的图像。最近一段时间，该公司已经成为AI画图领域最热门的平台。Midjourney的创始人兼首席执行官DavidHolz在一篇谷歌博客文章中说，“我们很自豪能与谷歌云合作，为我们的创意社区提供无缝体验，这些社区由谷歌在全球范围内提供支持可扩展的基础设施，从使用JAX在最新的v4TPU上训练我们算法的第四个版本，到在GPU上运行推理，我们对TPUv4允许我们的用户将他们充满活力的想法变为现实的速度印象深刻。”

谷歌在论文中表示，TPUv4比当代DSA芯片速度更快、功耗更低，如果考虑到互连技术，功率边缘可能会更大。能源效率提高了大约2-3倍，并且在公司称之为典型的本地数据中心的DSA上减少了大约20倍的二氧化碳排放量。通过使用具有3D环面拓扑的3KTPUv4切片，与TPUv3相比，谷歌的超算也能让LLM的训练时间大大减少。

性能、可扩展性和可用性使TPUv4超级计算机成为LaMDA、MUM和PaLM等大型语言模型(LLM)的主要算力。这些功能使5400亿参数的PaLM模型在TPUv4超算上进行训练时，能够在50天内维持57.8%的峰值硬件浮点性能。

此外，TPUv4的性能比TPUv3高2.1倍，性能功耗比提高2.7倍。基于TPUv4的超级计算机拥有4096块芯片，整体速度提高了约10倍。对于同等大小的系统，其芯片比基于英伟达A100芯片的系统快1.2-1.7倍，功耗低1.3-1.9倍，比GraphcoreIPUBow快4.3-4.5倍。

英伟达A100芯片与第四代TPU同时上市，并被用于GPT-4的训练。不过英伟达目前的旗舰是H100芯片（于2022年宣布），谷歌并未与之进行对比，因为H100在谷歌芯片之后上市，并且采用了更先进的Hopper架构和4nm工艺（与TPUv4的7nm相比）。谷歌论文作者还解释说，他们选择不将TPUv4与更新的H100进行比较，因为谷歌更喜欢在部署并用于运行生产应用程序之后撰写有关技术的论文。

对此，英伟达发言人拒绝置评，但本周他们“巧合”地发布了一篇博客，其中英伟达创始人兼首席执行官黄仁勋指出A100于三年前首次亮相，并且Nvidia最近的H100(Hopper)GPU提供的性能比基于MLPerf3.0基准的A100高出4倍。

(来源：Nvidia)

其实谷歌的TPUv4也是在三年前——也就是2020年服役，此后不断完善。“TPUv4s和A100s都在2020年部署，并且都使用7nm技术，”谷歌论文作者解释道。“更新的700WH100在2022年无法在AWS、Azure或GoogleCloud上使用。合适的H100匹配将是在类似时间框架和技术（例如，2023年和4nm）中部署的TPUv4的继任者。”

不过谷歌暗示正在研发一种新版本TPU，以期与英伟达H100芯片竞争，但没有提供细节。