谷歌 TPU 超级算力大模型已超英伟达

尽管目前尚无与 ChatGPT 匹敌的AI大型模型，但在计算能力方面，领导者可能并非微软和 OpenAI。目前，谷歌发布了其基于 TPU 的超级计算机在训练大型语言模型方面的详细信息，该超级计算机在速度和能效方面已经超过了英伟达的同类产品。



谷歌张量处理器（Tensor Processing Unit，TPU）是该公司为机器学习定制的专用集成电路（Application-Specific Integrated Circuit，ASIC）。自2016年发布第一代以来，TPU 便成为了 AlphaGo 背后的算力。与 GPU 相比，TPU 采用低精度计算，在几乎不影响深度学习处理效果的前提下显著降低了功耗并加快了运算速度。此外，TPU采用了脉动阵列等设计来优化矩阵乘法和卷积运算。

目前，谷歌90%以上的人工智能训练任务都在使用这些芯片，而TPU支持了包括搜索在内的谷歌的主要业务。图灵奖获得者、计算机架构领域的杰出人物大卫・帕特森（David Patterson）在2016年从 UC Berkeley 退休后，以杰出工程师的身份加入了谷歌大脑团队，为几代 TPU 的研发做出了突出贡献。



如今，TPU 已发展至第四代。4月11日，由 Norman Jouppi、大卫・帕特森等人发表的论文《TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings》详细阐述了谷歌如何利用自研光通信器件将 4000 多个芯片并联组成超级计算机，从而提升整体效能。

相较于 TPU v3，TPU v4 性能提升了2.1倍，功耗性能比增加了2.7倍。基于 TPU v4 的超级计算机拥有 4096 个芯片，整体速度提升约10倍。与类似规模的系统相比，谷歌的 TPU v4 比 Graphcore IPU Bow 快4.3-4.5倍，比 Nvidia A100 快1.2-1.7倍，功耗降低1.3-1.9倍。

除了芯片本身的计算能力外，芯片间互连已成为构建AI超级计算机的企业间竞争的关键因素。近期，大型语言模型（LLM）如谷歌的 Bard 和 OpenAI 的 ChatGPT 的规模正呈爆炸式增长，计算能力已经成为显著的瓶颈。

由于大型模型通常具有千亿级的参数量，数千个芯片必须共同分担训练任务，且训练过程可能持续数周甚至更长时间。谷歌的PaLM模型——迄今为止该公司披露的最大规模语言模型——在训练过程中被分配到两台拥有4000个TPU芯片的超级计算机上，耗时50天。

谷歌表示，通过光电路交换机（Optical Circuit Switch, OCS），其超级计算机可以轻松地动态重新配置芯片间的连接，有助于避免问题并实时调整以提高性能。

下图展示了 TPU v4 4×3 配置中的6个“面”的连接方式。每个面包含16条链路，每个模块总共有96条光链路连接到 OCS。为了提供3D环面的环绕连接，相对侧的连接必须接入相同的 OCS。因此，每个4×3模块的TPU连接到6×16÷2=48个 OCS。Palomar OCS 具有136×136规格（128个端口加上8个用于链路测试和修复的备用端口），因此来自64个4×3模块（每个包含64个芯片）的48对电缆共连接了48个 OCS，总计并联4096个 TPU v4 芯片。



根据这样的排布，TPU v4（中间的 ASIC 加上 4 个 HBM 堆栈）和带有 4 个液冷封装的印刷电路板 (PCB)。该板的前面板有 4 个顶部 PCIe 连接器和 16 个底部 OSFP 连接器，用于托盘间 ICI 链接。随后，八个 64 芯片机架构成一台 4096 芯片超算。



在超级计算机中，工作负载由不同规模的算力承担，称为切片，如64芯片、128芯片、256芯片等。下图展示了在主机可用性从99.0%到99.9%不等，以及在没有 OCS 时，不同切片大小的“有效输出”。若无 OCS，主机可用性必须达到99.9%才能提供合理的切片吞吐量。对于大多数切片大小，OCS 在99.0%和99.5%的可用性下也表现出良好的输出。



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

谷歌研究员 Norm Jouppi 和谷歌杰出工程师大卫・帕特森在一篇关于该系统的博客文章中写道：这种切换机制使得绕过故障组件变得容易，其灵活性甚至允许我们改变超级计算机互连的拓扑结构，以加速机器学习模型的性能。

TPU v4 比当代 DSA 芯片速度更快、功耗更低，如果考虑到互连技术，功率边缘可能会更大。通过使用具有 3D 环面拓扑的 3K TPU v4 切片，与 TPU v3 相比，谷歌的超算也能让 LLM 的训练时间大大减少。下图展示了谷歌自用的推荐模型（DLRM0）在不同芯片上的效率。TPU v3 比 CPU 快 9.8 倍。TPU v4 比 TPU v3 高 3.1 倍，比 CPU 高 30.1 倍。



性能、可扩展性和可用性使 TPU v4 超级计算机成为 LaMDA、MUM 和 PaLM 等大型语言模型 (LLM) 的主要算力。这些功能使 5400 亿参数的 PaLM 模型在 TPU v4 超算上进行训练时，能够在 50 天内维持 57.8% 的峰值硬件浮点性能。谷歌表示，其已经部署了数十台 TPU v4 超级计算机，供内部使用和外部通过谷歌云使用。