【论文评析】TPU v4：用于机器学习的光学可重构超级计算机，具有嵌入的硬件支持

题目：TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

名称：TPU v4：用于机器学习的光学可重构超级计算机，具有嵌入的硬件支持

论文：https://dl.acm.org/doi/10.1145/3579371.3589350

单位：谷歌

会议：ISCA 2023

TPU(Tensor Processing Unit)，张量处理单元，谷歌为机器学习而定制的ASIC（Application Specific Integrated Circuit，专用集成电路）芯片，可以实现在深度学习上比CPU和GPU更高的性能和能效比。

1，Introduction

机器学习 (ML) 模型的发展带来规模和算法方面的挑战。规模：大型语言模型（large language models，LLM）；算法方面：深度学习推荐模型（deep learning recommendation models，DLRM）

超算规模扩大：TPU v2（256）——TPU v4（4096），——>可靠性问题

TPU v4应对挑战的三个特性：

• 引入OCSes解决规模和可靠性障碍（允许 4K 节点超级计算机通过重新配置来容忍 1K CPU 主机在 0.1%–1.0% 的时间内不可用）
• 公开了对 DLRM（SparseCore 或 SC）中嵌入的硬件支持，DLRM 自 TPU v2 以来一直是 TPU 的一部分
• 满足all-to-all需求，使用OCSes 实现灵活的拓扑配置，包括具有更好二分特性的扭曲环面（twisted torus）

论文贡献：

• 描述并评估了超算中 OCSes 的首次生产部署，以及第一个允许拓扑重新配置以提高性能的部署
• 描述并评估了商业机器学习系统中第一个支持嵌入的加速器
• 记录了自 2016 年以来 ML 领域生产模型类型的变化，如下表
• 展示了 Google 如何使用 ML 来共同优化深度神经网络（DNN ）模型、OCS 拓扑和稀疏核（SparseCore）

2，可重构光交换机

2D 环面拓扑的一些环绕链路非常长，由于电互连的范围限制，它们必须是光学链路。光链路比电链路贵 10 倍以上。使用 3D 环面可增加对分带宽，OCS 的作用就像插板，可以跳过故障单元。

2.1 光电路交换

Google Palomar OCS 基于可在几毫秒内切换的 3D MEMS镜子。使用循环器在光纤中双向发送光，所需的端口和电缆数量减半。

电网络块的尺寸：每个 CPU 主机有 4 个 TPU v4，64 个 TPU v4 芯片装进 16 个 CPU 主机可以轻松装入一个机架。所以电网络块的尺寸为4*4*4

2.2 TPU v4 超级计算机的构建

每个面伸出16个链路，一个块共96个光链路；需要96/2=48个光交换机；

Palomar OCS 为 136×136（128 个端口加上 8 个用于链路测试和维修的备用端口）

每个节点视为一个4*4*4的cube，则每个OCS连接64+64个链路，4*4*4*（4*4*4）=4096

（64个机架其中8个）

2.3 OCS 容错优势

如果使用常规的固定互联架构，那么一个芯片出故障可能会影响整个系统工作。而在有了可重配置的光互连之后，需要做的只需要把出故障的芯片绕过，就不会影响整个系统的工作。

横坐标是工作负载的芯片规模，纵坐标是有效吞吐，不同颜色的线代表从99.0%到99.9%的芯片可靠率

假设芯片可靠率在99%的情况下，其整体系统的平均性能提升比不使用OCS可高达6倍

2.4 OCS部署优势

OCS 缩短了部署时间。在安装并测试所有 1024 芯片和所有电缆之前，TPU v3 系统无法使用。任何组件的交付延迟都会影响整个超级计算机的运行。对于 TPU v4，OCS 使每个机架独立，每个4^3 的cube块可以快速投入使用，提高了 TPU v4 超级计算机的成本效益。

2.5 OCS 调度简化优势

OCS 调度的优点 OCS 还简化了调度，从而提高了利用率。对于 TPU v3，256 个芯片片意味着调度程序必须找到 256 个连续的空闲芯片。对于 TPU v4，它可以从超级计算机的任何位置选择四个 43 块。切片甚至不需要是 2 的幂；它们可以是 4i×4j×4k，其中 0 < i ≤ j ≤ k。例如，用户可以请求几何尺寸为 4×4×12 的 192 TPU v4 切片。

2.6 OCS 模块化和安全优势

OCS 可以在几毫秒内切换电路，TPU v4 可以轻松更改拓扑以匹配应用程序、节点数量以及运行这些作业的系统。

OCS 还可以在不同切片之间实现气隙网络隔离，从而增强共享 TPU v4 超级计算机的多个客户的安全性。

2.7 定制OCS拓扑提高性能

可缩短 DNN 训练时间的三种基本并行类型：

1. 数据并行：每个芯片计算示例子集的前向和后向传递，并发送为其子集计算的梯度到其他芯片。

2. 模型（或张量）并行性：大型张量运算及其权重被划分到多个芯片上，以便每个芯片同时计算张量运算的子集。

3. 管道并行性：对于具有多个层的 DNN，每个芯片计算层的子集，并将层结果传达给持有相邻层的芯片。

用户可以更改 TPU v4 拓扑以匹配所使用的并行类型，例如，对于 512 切片，管道并行可能需要雪茄形状 (4×4×32)，而不是传统的 8^3 立方体 (8×8×8) 。对于嵌入重型应用程序通常需要的最高二等分带宽，传统的 8^3 立方体是首选。

2.8 扭曲torus

TPU v4 节点的 4×2 切片的常规（顶部）和扭曲环面（底部）拓扑示例。

电连接（红色虚线）保持固定。通过利用 OCS 的灵活性，光连接（蓝色实线）可以从矩形环面重新配置为扭曲环面，而无需对机器进行任何物理重新布线，唯一的变化是路由表。

这么做的意义在于：将非完美切片重构为完美立方体切片，沿每个维度具有相同数量 TPU 的对称环面可最大限度地减少延迟并最大化二等分带宽。

与常规环面相比，扭曲环面在 4×4×8 和 4×8×8 切片上的整体吞吐量分别提高了 1.63 倍和 1.31 倍

2.9 OCS 成本

OCS 成本小于 TPU v4 超级计算机总成本的 5%，且小于总功耗的 3%（功率和成本核算包括整个光学结构，包括光学模块、光纤和 OCS 基础设施）

课程报告PPT