本论文是 Google 发表在Proceedings of the 50th Annual International Symposium on Computer Architecture ISCA ‘23 2023上的Industry Section工作。

Google 在 ASIC (Application Specific Integrated Circuit, 专用集成芯片) 上的工作始于 2015 年第一代 TPU (Tensor Processing Unit), 在2017、2018、2021年分别推出了 v2, v3, v4 版本，核心数量从 64 个以 4 倍增长的速度达到 4096 个，此后在 2023 年推出 TPU v5e/v5p，2024年推出 TPU v6，具备 1836 TOPS¹，但是还是低于英伟达 H100 的 3026 TOPS （高于 A100 的 624 TOPS²），尽管算力具备差距，其基于 Palomar OCS (Optical Communcation Switch, 光交换机) 的网络提高了计算集群的性能。

TPU v4 是谷歌首次引入 OCS 优化通信，并提出 Sparse Core 优化稀疏计算的版本，下面将重点介绍这两方面。

Entry：Zotero link URL link

稀疏计算核

Embedding 层的计算

Embedding 处理离散型分类特征（Categorical Features），是稀疏化的典型计算范式，NLP/搜推算法仅支持字符串形式输入是标示为离散的稀疏向量特征，因此 Embedding 过程不适合映射到硬件上的矩阵乘法单元进行 Tensor 计算。

同时，将这些稀疏变量放在 Tensor Cores 上进行计算会频繁发生小规模的内存访问动作，很容易会触及到 CPU 和 DRAM 的性能瓶颈，尤其是在 TPU 和 CPU 4:1 的 TPU 集群上，而在此之上，网络中心的延迟也会进一步加剧这种负面影响。因此，能够高效处理稀疏变量的计算范式就被应用到了 TPU 之上。

深度学习中由于神经网络通常在稠密 Tensor 上计算性能更优，因此就需要对稀疏的矩阵进行压缩，通过特定的 Embedding 层去生成一个稠密的原数据表达，这通常是 NLP 或者搜推算法的第一层，例如 Deep & Wide Model 中，用户的稀疏特征被先被处理成稠密 Embeddings 再进入中间的隐藏层。

Sparse Core 核心架构

下图是 Sparse Core 的架构示意图，图中的深蓝色方框就是 SC 中最通用的单元 - 16 个 tiles. 每个 tile 都有一个关联的 HBM 通道，并支持多个内存访问。每个 tile 都有一个 Fetch Unit、一个可编程的 8 路 SIMD 向量处理单元（这里是 scVPU，不要与 TPU v4 中的 Tensor Core 的 VPU 混淆）、以及一个 Flush Unit。Fetch Unit，从 HBM 读取激活函数和参数到 tiles 的 2.5 MiB 稀疏向量内存（Spmem）的切片中。Flush Unit 在反向传播过程中将更新的参数写入 HBM。此外，五个跨通道单元（图中的金色方框）在 16 个 Spmem 上执行与他们名称相对应的嵌入操作。与 TPU v1 一样，这些单元执行类似 CISC 的指令，并处理可变长度的输入，其中每条指令的运行时间取决于数据。

TPU v4 Pod

TPU v4 Pod 提供了 exaflops 的算力，原因就是因为其强大的芯片间互联能力。谷歌将 $4×4×4=64$ 个 TPU v4 芯片互联在一起，形成一个立方体结构（Cube）。再把 $4×4×4$ Cube 用光互联连在一起形成一个总共有 4096 个 TPU v4 超级计算机。

光互联

谷歌在 TPU v4 Pod 中应用了光交换器（OCS）为了避免计算等通信。OCS 由 64 （4×4×4）颗 TPU 构成一组 Slice 间互连，实现了 Pod 内 Slice 间全光互连（4096 TPUs）；当然也可用于 Pod 之间互连而不仅局限于芯片之间。

在使用可配置光互连（以及光路开关）时，假设芯片可靠率在 99%的情况下，其整体系统的平均性能提升比不使 OCS 可高达 6 倍。

拓扑结构

谷歌在 TPU v4 Pod 中采取了 3D Torus 的拓扑结构实现芯片之间的互联。在 3D Torus 网络中，每个节点通过三个维度的连接与相邻节点相连，形成一个三维的网格结构。这种高度互联的结构允许芯片在 X-Y-Z 三个维度上形成一个连续循环，提供高带宽和低延迟的通信能力，非常适合于高性能计算和大规模并行处理。TPU v4 的每个端口对应一个 TPU v4 芯片，端口连接交换机提供 6 Tb/sec 的带宽，作为网络接口卡和 3D 环面网络的基础。这种设计确保了 TPU v4 在处理大规模并行计算任务时的高效性和稳定性。

下图是一个 3D Torus Cube 的图示，我们可以看到每一个节点都会和上下、左右、前后的节点互联，而处于这个立方体边上的芯片会连接到在它相对面上的节点从而实现每一个节点都与 6 个节点互联。 Cube 如果要实现 6 面连接，每个面需要 16 条链路，每个 Cube 总共有 96 条光链路连接到 OCS 上。为了提供整个 3D Torus 的连接，相对的侧面需要连接到相同的 OCS 上面（图中不同色块和 Cube 之间的长环）。因此，每个 Cube 会连接到 $6×16÷4=48$ 个 OCS 上。这 48 个 OCS 一共连接来自 64 个 Cube 的 48 对光缆，总共并联 4096 个 TPU v4 芯片。

对比英伟达 DGX SuperPod，其第四代 NVLink/NVSwitch 最多可以连接 32 个 node 总 256 颗 H100 芯片，并实现每颗 GPU 900G/s 的互连带宽。NV 每机架 4 台 DGX（共 32 颗 H100 GPU），机架内/外需要光纤连接。NV 的每机架算力密度相对更小/更窄，且需要更多的收发激光器和光纤线材，网络成本高。如果 NV 部署 4096 颗 GPU 集群，必须切分成更多个 SuperPod 并独立规划互连网络层，中间完成多层交换，集群内总计需要采购大约 568 个 Infiniband Switch。

反观 TPU v4 Pod，与超级计算机一样，工作负载由不同规模的算力承担，称为切片：64 芯片、128 芯片、256 芯片等。与 Infiniband 相比，OCS 的成本更低、功耗更低、速度更快，成本不到系统成本的 5%，功率不到系统功率的 3%。每个 TPU v4 都包含 SparseCores 数据流处理器，可将依赖嵌入的模型加速 5 至 7 倍，但仅使用 5% 的裸片面积和功耗。

光路交换机

OCS 解决了传统交换机需要从电信号到光信号再到电信号的额外开销，其基本的构成是 Palomar芯片。

光路开关芯片

光路有多开关芯片叫做 Palomar，使用的是基于 MEMS 反射镜阵列的技术，具体原理是使用一个 2D MEMS 反射镜阵列，通过控制反射镜的位置来调整光路，从而实现光路切换。使用 MEMS 光路开关芯片可以实现低损耗，低切换延迟（毫秒级别）、低功耗、低成本。

2D MEMS 阵列

每个MEMS 反射镜阵列陶瓷封装内部是单个大型芯粒，芯粒里面有 176 个可单独控制的微反射镜。其热成像图中每一个深色的圆形都代表了一个微反射镜，每一个微反射镜四周都有四个梳状驱动区域，用于在两个方向上旋转反射镜，使得光可以进行光切换。

下面这个图解释了 Palomar OCS 的工作原理。绿色线条表示带内光信号路径。与带内信号路径叠加的是一个 850 纳米波长的监控通道（红色箭头），用于校准镜面。输入/输出的光学信号通过二维（2D）光纤准直器阵列（Fiber collimator array）进入光学核心，每个准直器阵列由一个 $N \times N$ 光纤阵列和 2D 透镜阵列组成。这个的核心是其中的两组 2D MEMS 阵列。MEMS 的镜面被在监控通道的校准下，被驱动、倾斜以将信号切换到相应的输入/输出准直光纤。

TPUv4 优缺点分析

优点

低时延：3D Tours 因其相邻节点之间的短而直接的连线，可以换来更低的延迟；尤其当节点间需要运行密集 I/O 的、紧耦合的并行任务时特别有用。
低网络代价：对于相同数量的节点，3D Torus 拓扑网络直径低于 Clos 拓扑，两者相比之下，前者的交换机/线材/连接器的保有量更低，网络层次更少，节省硬件成本。
路由可重配：谷歌 OCS 网络支持动态可重配路由，Silic 集群在部署之后可以立即投入生产，无需等待整个网络收敛；并且这种特性更容易隔离/下线故障节点。
更好集群布局：集群布局让物理连接上相临的节点间在逻辑上临近，让密集 I/O 通信、data-flow 发生在局部流域，换来更低的通信开销；同时优化了延迟和功耗。这就是 3D Torus 将大集群逻辑切割成紧耦合的局部域，局部互连并共享作业。

缺点

系统成熟度低：Clos 拓扑本身具备非阻塞特点，性能能够始终保持一致且可以预测，其所有输入/输出都是全带宽同时连接，无冲突无阻塞，这在 3D Tours 拓扑中无法保证。
拓扑僵硬：在 Clos 这种 Spine-Leaf 脊叶拓扑中，扩容新的叶交换机相对简单，无需更改当前架构；相比之下，扩缩 3D Tours 结构比较复杂和耗时，可能需要重新配置整个拓扑。
负载均衡问题：Clos 网络在任意两个节点之间提供更多路径，从而实现负载均衡和冗余；虽然 3D Tours 结构也提供多路径冗余，但显而易见 Clos 的替代路径数量更多，具体取决于网络的配置。

Last modified on 2024-12-23