本周二，谷歌公布了其演习措辞大模型的超级打算机的细节，基于 TPU 的超算系统已经可以比英伟达的同类更加快速、节能。
谷歌张量处理器（tensor processing unit，TPU）是该公司为机器学习定制的专用芯片（ASIC），第一代发布于 2016 年，成为了 AlphaGo 背后的算力。与 GPU 比较，TPU采取低精度打算，在险些不影响深度学习处理效果的条件下大幅降落了功耗、加快运算速率。同时，TPU 利用了脉动阵列等设计来优化矩阵乘法与卷积运算。
当前，谷歌 90% 以上的人工智能演习事情都在利用这些芯片，TPU 支撑了包括搜索的谷歌紧张业务。作为图灵奖得主、打算机架构巨擘，大卫・帕特森（David Patterson）在 2016 年从 UC Berkeley 退休后，以精彩工程师的身份加入了谷歌大脑团队，为几代 TPU 的研发做出了卓越贡献。
如今 TPU 已经发展到了第四代，谷歌本周二由 Norman Jouppi、大卫・帕特森等人揭橥的论文《 TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings 》详细先容了自研的光通信器件是如何将 4000 多块芯片并联成为超级打算机，以提升整体效率的。
TPU v4 的性能比 TPU v3 高 2.1 倍，性能功耗比提高 2.7 倍。基于 TPU v4 的超级打算机拥有 4096 块芯片，整体速率提高了约 10 倍。对付类似大小的系统，谷歌能做到比 Graphcore IPU Bow 快 4.3-4.5 倍，比 Nvidia A100 快 1.2-1.7 倍，功耗低 1.3-1.9 倍。
除了芯片本身的算力，芯片间互联已成为构建 AI 超算的公司之间竞争的关键点，最近一段韶光，谷歌的 Bard、OpenAI 的 ChatGPT 这样的大措辞模型（LLM）规模正在爆炸式增长，算力已经成为明显的瓶颈。
由于大模型动辄千亿的参数量，它们必须由数千块芯片共同分担，并持续数周或更永劫光进行演习。谷歌的 PaLM 模型 —— 其迄今为止最大的公开表露的措辞模型 —— 在演习时被拆分到了两个拥有 4000 块 TPU 芯片的超级打算机上，用时 50 天。
谷歌表示，通过光电路交流机（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 个用于链路测试和修复的备用端口），因此 48 个 OCS 连接来自 64 个 4×3 块（每个 64 个芯片）的 48 对电缆，统共并联 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 和谷歌精彩工程师大卫・帕特森在一篇关于该系统的博客文章中写道。「这种灵巧性乃至许可我们改变超级打算机互连的拓扑构造，以加速机器学习模型的性能。」
在新论文上，谷歌着重先容了稀疏核（SparseCore，SC）的设计。在大模型的演习阶段，embedding 可以放在 TensorCore 或超级打算机的主机 CPU 上处理。TensorCore 具有宽 VPU 和矩阵单元，并针对密集操作进行了优化。由于小的聚拢 / 分散内存访问和可变长度数据交流，在 TensorCore 上放置嵌入实在并不是最佳选择。在超级打算机的主机 CPU 上放置嵌入会在 CPU DRAM 接口上引发阿姆达尔定律瓶颈，并通过 4:1 TPU v4 与 CPU 主机比率放大。数据中央网络的尾部延迟和带脱期制将进一步限定演习系统。
对此，谷歌认为可以利用 TPU 超算的总 HBM 容量优化性能，加入专用 ICI 网络，并供应快速网络 / 分散内存访问支持。这导致了 SparseCore 的协同设计。
SC 是一种用于嵌入演习的特定领域架构，从 TPU v2 开始，后来在 TPU v3 和 TPU v4 中得到改进。SC 相对划算，只有芯片面积的约 5% 和功率的 5% 旁边。SC 结合超算规模的 HBM 和 ICI 来创建一个平坦的、全局可寻址的内存空间（TPU v4 中为 128 TiB）。与密集演习中大参数张量的全部归约比较，较小嵌入向量的全部传输利用 HBM 和 ICI 以及更细粒度的分散 / 聚拢访问模式。
作为独立的核心，SC 许可跨密集打算、SC 和 ICI 通信进行并行化。下图显示了 SC 框图，谷歌将其视为「数据流」架构（dataflow），由于数据从内存流向各种直接连接的专用打算单元。
最通用的 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 的指令并对可变长度输入进行操作，个中每条指令的运行韶光都取决于数据。
在特定芯片数量下，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 倍改进。
虽然谷歌直到现在才公布有关其超级打算机的详细信息，但自 2020 年以来，基于 TPU 的 AI 超算一贯在位于俄克拉荷马州的数据中央发挥浸染。谷歌表示，Midjourney 一贯在利用该系统演习其模型，最近一段韶光，后者已经成为 AI 画图领域最热门的平台。
谷歌在论文中表示，对付同等大小的系统，其芯片比基于英伟达 A100 芯片的系统快 1.7 倍，能效高 1.9 倍，后者与第四代 TPU 同时上市，并被用于 GPT-4 的演习。
对此，英伟达发言人谢绝置评。
当前英伟达的 AI 芯片已经进入 Hopper 架构的时期。谷歌表示，未对第四代 TPU 与英伟达目前的旗舰 H100 芯片进行比较，由于 H100 在谷歌芯片之后上市，并且采取了更前辈的制程。
但同样在此，谷歌暗示了下一代 TPU 的操持，其没有供应更多细节。Jouppi 见告路透社，谷歌拥有开拓「未来芯片的康健管道」。
TPU v4 比当代 DSA 芯片速率更快、功耗更低，如果考虑到互连技能，功率边缘可能会更大。通过利用具有 3D 环面拓扑的 3K TPU v4 切片，与 TPU v3 比较，谷歌的超算也能让 LLM 的演习韶光大大减少。
性能、可扩展性和可用性使 TPU v4 超级打算机成为 LaMDA、MUM 和 PaLM 等大型措辞模型 (LLM) 的紧张算力。这些功能使 5400 亿参数的 PaLM 模型在 TPU v4 超算上进行演习时，能够在 50 天内坚持 57.8% 的峰值硬件浮点性能。
谷歌表示，其已经支配了数十台 TPU v4 超级打算机，供内部利用和外部通过谷歌云利用。
本文作者：泽南，来源：机器之心，原文标题：《谷歌TPU超算，大模型性能超英伟达，已支配数十台：图灵奖得主新作》
本文来自华尔街见闻，欢迎下载APP查看更多
