谷歌近日公布了其用于训练人工智能（AI）模型的超级计算机的最新细节。该公司称，这些系统比英伟达公司的同期系统更快、更节能。随着OpenAI的ChatGPT、谷歌的Bard等AI聊天机器人的竞争愈演愈烈，改善芯片间的连接已成为开发AI超算的科技公司的一个关键竞争点……
当地时间周二（4月4日），Alphabet旗下的谷歌公司公布了其用于训练人工智能（AI）模型的超级计算机的最新细节。该公司称，这些系统比英伟达公司的同期系统更快、更节能。
据了解，谷歌张量处理器（tensor processing unit，简称TPU芯片）是该公司为机器学习定制的专用芯片（ASIC），第一代发布于2016年，为当时大热的AlphaGo提供算力。与GPU相比，TPU采用低精度计算，在几乎不影响深度学习处理效果的前提下大幅降低了功耗、加快运算速度。同时，TPU使用了脉动阵列等设计来优化矩阵乘法与卷积运算。
Google CEO Sundar Pichai 在 Google I/O 2021 上宣布 TPU v4。（来源：Google）
作为图灵奖得主、计算机架构巨擘，大卫・帕特森（David Patterson）在 2016 年从 UC Berkeley 退休后，以杰出工程师的身份加入了谷歌大脑团队，为几代 TPU 的研发做出了卓越贡献。
谷歌表示，目前公司90%以上的AI训练工作都是通过谷歌自主研制的TPU芯片完成的。其AI训练的过程是：通过给模型输入数据，使其能够用类似人类的文本回应、并具备生成图像等功能。
当前谷歌使用的是最新的第四代TPU芯片（TPU v4），谷歌公司周二发布的一篇题为《 TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings 》的论文，详细阐述了该公司是如何利用自己定制的光路交换器（OCS），将4000多个芯片组合到一台超级计算机中，从而帮助连接一台台独立机器。谷歌表示，其已经部署了数十台采用第四代TPU打造的超级计算机，供内部使用和外部通过谷歌云使用。
相关论文
改善芯片间的连接成AI超算关键竞争点
随着OpenAI的ChatGPT、谷歌的Bard等AI聊天机器人的竞争愈演愈烈，改善芯片间的连接已成为开发AI超算的科技公司的一个关键竞争点。
由于ChatGPT、Bard等产品背后使用的大型语言模型的规模太庞大了，动辄千亿的参数量无法存储在单一芯片上。所以，这些模型必须分散在数千个芯片上，然后这些芯片必须同时工作，花费数周或更长时间来训练语言模型。
谷歌的PaLM是该公司迄今为止公开披露的最大的语言模型。其训练方式是：通过将其分配到两台各包含4000个芯片的超级计算机中，历时50天进行训练。
谷歌周二称，其通过光电路交换机，其超级计算机可以轻松地动态重新配置芯片之间的连接，有助于避免出现问题并实时调整以提高性能。
下图展示了 TPU v4 4×3 方式 6 个「面」的链接。每个面有 16 条链路，每个块总共有 96 条光链路连接到 OCS 上。要提供 3D 环面的环绕链接，相对侧的链接必须连接到相同的 OCS。因此，每个 4×3 块 TPU 连接到 6 × 16 ÷ 2 = 48 个 OCS 上。Palomar OCS 为 136×136（128 个端口加上 8 个用于链路测试和修复的备用端口），因此 48 个 OCS 连接来自 64 个 4×3 块（每个 64 个芯片）的 48 对电缆，总共并联 4096 个 TPU v4 芯片。
(来源：Google)
根据这样的排布，TPU v4（中间的 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%。每个 TPU v4 都包含 SparseCores 数据流处理器，可将依赖嵌入的模型加速 5 至 7 倍，但仅使用 5% 的裸片面积和功耗。
该公司的研究员Norm Jouppi和工程师David Patterson在上述论文中写道，“电路切换可以很容易地绕过故障组件…这种灵活性甚至允许我们改变超级计算机互连的拓扑结构，以加速机器学习模型的性能。”
稀疏核（SparseCore，SC）设计
在新论文上，谷歌着重介绍了稀疏核（SparseCore，SC）的设计。在大模型的训练阶段，嵌入训练可以放在 TensorCore 或超级计算机的主机 CPU 上处理。TensorCore 具有宽 VPU 和矩阵单元，并针对密集操作进行了优化。由于小的聚集 / 分散内存访问和可变长度数据交换，在 TensorCore 上放置嵌入其实并不是最佳选择。在超级计算机的主机 CPU 上放置嵌入会在 CPU DRAM 接口上引发阿姆达尔定律瓶颈，并通过 4:1 TPU v4 与 CPU 主机比率放大。数据中心网络的尾部延迟和带宽限制将进一步限制训练系统。
对此，谷歌认为可以使用 TPU 超算的总 HBM 容量优化性能，加入专用 ICI 网络，并提供快速收集 / 分散内存访问支持。这导致了 SparseCore 的协同设计。
SC 是一种用于嵌入训练的特定领域架构(DSA)，从 TPU v2 开始，后来在 TPU v3 和 TPU v4 中得到改进。SC 相对划算，只使用芯片面积和功率的约 5% 左右。SC 结合超算规模的 HBM 和 ICI 来创建一个平坦的、全局可寻址的内存空间（TPU v4 中为 128 TiB）。与密集训练中大参数张量的全部归约相比，较小嵌入向量的全部传输使用 HBM 和 ICI 以及更细粒度的分散 / 聚集访问模式。
作为独立的核心，SC 允许跨密集计算、SC 和 ICI 通信进行并行化。下图显示了 SC 框图，谷歌将其视为「数据流」架构（dataflow），因为数据从内存流向各种直接连接的专用计算单元。
(来源：Google)
最通用的 SC 单元是 16 个计算块（深蓝色框）。每个 tile 都有一个关联的 HBM 通道，并支持多个未完成的内存访问。每个 tile 都有一个 Fetch Unit、一个可编程的 8-wide SIMD Vector Processing Unit 和一个 Flush Unit。获取单元将 HBM 中的激活和参数读取到 2.5 MiB 稀疏向量内存 (Spmem) 的图块切片中。scVPU 使用与 TC 的 VPU 相同的 ALU。Flush Unit 在向后传递期间将更新的参数写入 HBM。此外，五个跨通道单元（金色框）执行特定的嵌入操作，正如它们的名称所解释的那样。
与 TPU v1 一样，这些单元执行类似 CISC 的指令并对可变长度输入进行操作，其中每条指令的运行时间都取决于数据。
(来源：Google)
在特定芯片数量下，TPU v3/v4 对分带宽比高 2-4 倍，嵌入速度可以提高 1.1-2.0 倍。
下图展示了谷歌自用的推荐模型（DLRM0）在不同芯片上的效率。TPU v3 比 CPU 快 9.8 倍。TPU v4 比 TPU v3 高 3.1 倍，比 CPU 高 30.1 倍。 谷歌探索了 TPU v4 超算用于 GPT-3 大语言模型时的性能，展示了预训练阶段专家设计的 1.2 倍改进。
(来源：Google)
新版本TPU与英伟达A100对比
虽然谷歌直到现在才披露关于其超级计算机的详细信息，不过事实上，自2020年以来，该系统已经在公司内部上线，应用于该公司位于俄克拉荷马州的一个数据中心。
来自位于俄克拉荷马州的谷歌 ML 集群的 TPU v4 pod 的八分之一，该公司声称该集群使用约 90% 的无碳能源运行。 （来源：Google）
据悉，初创公司Midjourney使用TPU v4系统来训练其模型，在输入几句文字后，模型就会生成新的图像。最近一段时间，该公司已经成为 AI 画图领域最热门的平台。Midjourney 的创始人兼首席执行官 David Holz 在一篇谷歌博客文章中说，“我们很自豪能与谷歌云合作，为我们的创意社区提供无缝体验，这些社区由谷歌在全球范围内提供支持 可扩展的基础设施，从使用 JAX 在最新的 v4 TPU 上训练我们算法的第四个版本，到在 GPU 上运行推理，我们对 TPU v4 允许我们的用户将他们充满活力的想法变为现实的速度印象深刻。”
谷歌在论文中表示，TPU v4 比当代 DSA 芯片速度更快、功耗更低，如果考虑到互连技术，功率边缘可能会更大。能源效率提高了大约 2-3 倍，并且在公司称之为典型的本地数据中心的 DSA 上减少了大约 20 倍的二氧化碳排放量。通过使用具有 3D 环面拓扑的 3K TPU v4 切片，与 TPU v3 相比，谷歌的超算也能让 LLM 的训练时间大大减少。
性能、可扩展性和可用性使 TPU v4 超级计算机成为 LaMDA、MUM 和 PaLM 等大型语言模型 (LLM) 的主要算力。这些功能使 5400 亿参数的 PaLM 模型在 TPU v4 超算上进行训练时，能够在 50 天内维持 57.8% 的峰值硬件浮点性能。
此外，TPU v4 的性能比 TPU v3 高 2.1 倍，性能功耗比提高 2.7 倍。基于 TPU v4 的超级计算机拥有 4096 块芯片，整体速度提高了约 10 倍。对于同等大小的系统，其芯片比基于英伟达A100芯片的系统快1.2-1.7倍，功耗低 1.3-1.9 倍，比 Graphcore IPU Bow 快 4.3-4.5 倍。
英伟达A100芯片与第四代TPU同时上市，并被用于GPT-4的训练。不过英伟达目前的旗舰是H100芯片（于 2022 年宣布），谷歌并未与之进行对比，因为H100在谷歌芯片之后上市，并且采用了更先进的Hopper 架构和4nm工艺（与 TPU v4 的 7nm 相比）。谷歌论文作者还解释说，他们选择不将 TPU v4 与更新的 H100进行比较，因为谷歌更喜欢在部署并用于运行生产应用程序之后撰写有关技术的论文。
对此，英伟达发言人拒绝置评，但本周他们“巧合”地发布了一篇博客，其中英伟达创始人兼首席执行官黄仁勋指出 A100 于三年前首次亮相，并且 Nvidia最近的 H100 (Hopper) GPU 提供的性能比基于 MLPerf 3.0 基准的 A100 高出 4 倍。
(来源：Nvidia)
其实谷歌的 TPU v4 也是在三年前——也就是 2020 年服役，此后不断完善。“TPU v4s 和 A100s 都在 2020 年部署，并且都使用 7nm 技术，”谷歌论文作者解释道。 “更新的 700W H100 在 2022 年无法在 AWS、Azure 或 Google Cloud 上使用。合适的 H100 匹配将是在类似时间框架和技术（例如，2023 年和 4nm）中部署的 TPU v4 的继任者。”
不过谷歌暗示正在研发一种新版本TPU，以期与英伟达H100芯片竞争，但没有提供细节。
您可能感兴趣
