moe

TPU Torus and OCS Optical Switching

TPU Torus and OCS Optical Switching

Posted by George Lin on June 21, 2026

引言:两条路的分岔——GPU 的交换机帝国与 TPU 的邻居网络

在 AI 训练芯片的互联战场上,NVIDIA 和 Google 各自选择了一条截然不同的路。

NVIDIA 的路是交换机之路。从 DGX-1 的 8 卡 PCIe Switch,到 DGX-2 的 NVSwitch 初代,再到 H100 的第三代 NVSwitch——每一代都在扩张交换机覆盖的 GPU 数量。NVIDIA 的终极愿景是一个”全互联”的世界:所有 GPU 都能以满带宽和均匀延迟与任何其他 GPU 通信,就像它们共处一台巨型计算机中。NVSwitch 就是这个愿景的执行者——一个非阻塞 crossbar,让每对 GPU 之间的通信都像直接连线一样。到了 NVL72,这个愿景扩展到 72 块 GPU,130 TB/s 的聚合全互联带宽。

Google 的路是邻居之路。从 TPU v2 开始,Google 就没有选择交换机。TPU 芯片之间直接通过 ICI(Inter-Chip Interconnect)链路相连,形成 2D 或 3D 的 Torus(环形网络)。每个芯片只连接物理上相邻的 4 个(2D Torus)或 6 个(3D Torus)邻居。数据从芯片 A 到远处的芯片 D,需要穿过中间芯片 B 和 C——每一跳都会消耗带宽。

这听起来像是一个明显的劣势。但当 Google 在 2020 年部署 TPU v4 时,他们在这条”邻居之路”上加了一个颠覆性的变量:光电路交换(OCS, Optical Circuit Switching)。48 个 Palomar MEMS 光交换机,以不到总成本 5%、不到总功耗 3% 的代价,让 4096 块 TPU v4 芯片的互联拓扑可以在毫秒级重新配置——从规则的 3D Torus 切换为扭曲 Torus(Twisted Torus),从 8×8×8 的立方体变形为 4×8×16 的”雪茄”,甚至可以在训练任务之间动态改变。

而结果是令人震惊的:这个看上去”简单粗暴”的邻居网络,在实际 ML 训练中的表现不仅不逊色,在能效和碳排放上甚至大幅领先同期竞品——1.2x-1.7x 快于 A100,功耗低 1.3x-1.9x,碳足迹低至竞品的 ~5%。

一条路靠增加交换机复杂度来隐藏拓扑,另一条路靠重新配置拓扑来适应流量。 这篇文章,我们要深入剖析第二路的全部技术细节:3D Torus 的物理构建、OCS 光学交换的革命性设计、动态拓扑重配置、嵌入加速的 SparseCore、以及让这一切在 99.98% 可用性下可靠运行的软件基础设施。


一、系统剖析:4096 芯片的 3D Torus 是如何搭起来的

1.1 从 TPU 芯片到 4x4x4 Cube

要理解 TPU v4 的互联架构,我们要从最小单元开始——单个 TPU v4 芯片。

每个 TPU v4 芯片包含 2 个 TensorCore(TC),每个 TC 内含 4 个 128×128 的矩阵乘法单元(MXU)和一个 128 通道的向量处理单元(VPU)。两个 TC 共享 128 MiB 的片上 CMEM(Common Memory)。HBM2 容量 32 GiB,带宽 1200 GB/s。芯片还包含 4 个 SparseCore——一个占了约 5% 芯片面积和功率、但能将嵌入操作加速 5-7 倍的专用数据流处理器。

在互联层面,每块 TPU v4 有 6 条 ICI 链路,每条带宽 50 GB/s(即 400 Gbit/s),分别指向 3D 空间的 ±X、±Y 和 ±Z 六个方向。注意这个带宽值得和竞品做直接对比——InfiniBand NDR200 的单链路带宽约 50 GB/s(200 Gbit/s),NVLink 每 link 约 25 GB/s(A100)——所以 ICI 的单链路带宽在同一个量级,但它用的是邻居直连而非交换机汇聚的模型。

4 个 TPU v4 芯片共享一个 CPU Host,装在同一个液冷托盘(tray)上。它们的 ICI 链路由 PCB 上的布线连接成一个 2×2 的 mesh。16 个这样的托盘组装成一个机架(rack)——64 块 TPU v4 芯片,通过机架内部的被动电缆构成一个 4×4×4 的 3D mesh。

但关键是环绕链路(wrap-around links)。4×4×4 mesh 如果不形成环,边缘芯片的带宽就损失了一半(它们只有 3 个或 4 个邻居,而非 6 个)。Torus 的解法是把对边的链路对接——4×4×4 的面向对面的方向延伸,在逻辑上闭合。这使得 64 芯片的集合从 3D mesh 变成 3D Torus(也叫做 4×4×4 Cube,或简写为 43 块)。

每个 4×4×4 Cube 有 6 个面,每个面有 4×4=16 条 ICI 链路需要伸出机架连接外部。 6 个面 × 16 条 = 96 条对外光链路。因为在 Torus 中,方向相反的两条链路(比如 +X 面的某条链路和 -X 面相同位置的那条链路)必须连接到同一个交换设备来实现闭合——所以每个 Cube 实际上需要连接到 96 ÷ 2 = 48 个 OCS

1.2 Palomar OCS:光电路交换机的物理原理

Google 的 Palomar OCS 是这个架构的核心创新。它的技术路线和电子交换机(如 InfiniBand 交换机或 NVSwitch)截然不同。

Palomar OCS 基于 3D MEMS(微机电系统)微镜阵列。它的核心是一个 136×136 的端口交换机,其中 128 个端口用于数据传输,8 个作为备用端口用于测试和维修。其工作原理可以简化为:

  1. 光纤从输入端传入,被准直器(collimator)转换为自由空间中的光束
  2. 光束经过第一个 2D MEMS 镜面阵列(136 个可单独控制的微镜),通过静电驱动器在两个方向上旋转微镜的角度
  3. 光束再经过第二个 2D MEMS 镜阵列,最终被引导到目标输出光纤
  4. 两个微镜阵列之间的路径完全是光学路径——没有光电转换、没有串行/解串(SerDes)、没有数据包处理

这意味着 OCS 不关心带宽是多少。 只要光能通过,不管是 400 Gbit/s 还是多波长复用的多个 Tbit/s,OCS 都照单全收——因为它只是在物理层反射光,不处理数据。这和电子交换机必须对每个数据包做解析、查表、转发的模型有着根本性的差异。

关键参数:

  • 切换时间:毫秒级(MEMS 机械运动需要物理偏转微镜,但这个时间对于拓扑重配置来说是完全可以接受的——拓扑在一个训练任务内保持静态,只在任务间或维护时切换)
  • 功耗:<3% 的全系统功耗——因为 MEMS 镜面保持固定角度所需的维持功率极小,远低于电子交换机对每个数据包进行路由决策的功耗
  • 成本:<5% 的全系统成本
  • 插入损耗:使用光环形器(circulator)在一根光纤中双向传输,将所需光纤和端口数量减半

1.3 从 Cube 到 4096 芯片超级计算机

48 个 Palomar OCS 连接 64 个 4×4×4 Cube,就构成了完整的 4096 芯片 TPU v4 超级计算机。具体来说:

  • 每个 Cube 的 6 个面 × 16 条链路 ÷ 2 = 48 对光纤入/出——每对一个 OCS
  • 48 个 OCS 各连接来自 64 个 Cube 的 128 个端口(每个 OCS 128 端口 + 8 个备用端口 = 136 端口总量)
  • 这意味着每个 OCS 恰好连接了 64 个不同的 Cube(每个 Cube 贡献一个”方向对”的链路)

物理布局是 8 排机架,每排 8 个机架,每个机架一个 43 Cube。电信号在机架内部走差分线,出机架时在光纤连接器上完成电-光转换,之后直到目标机架才再次做光-电转换——中间经过的路径全程是光:OCS 的输入端口 → MEMS 镜面反射 → OCS 的输出端口 → 目标机架。没有中间转发、没有缓存、没有路由查表。

这个架构的关键优势之一:增量部署。 在 TPU v3 时代,所有 1024 个芯片和所有电缆必须全部安装和测试完毕后,整个系统才可以投入使用。任何一个组件的延迟交付都会拖住整个超级计算机的后腿。但有了 OCS,每个 Cube 是独立的——一个 Cube 的 64 芯片和配套电缆安装测试完毕,就可以立刻投入生产运行。这种模块性不仅加快了交付速度,也显著提升了投资回报率。

1.4 电-光转换的物理路劲

一个数据包从 TPU v4 芯片 A 出发到达处于不同 Cube 中的目标芯片 B,经历的过程如下:

  1. ICI 链路层:数据从芯片 A 的 ICI 输出端口通过 PCB 走线和被动电缆到达机架的面板光纤连接器
  2. 电-光转换:在发货端的 OSFP 光纤模块中,电信号转换为光信号
  3. OCS 光路径:光信号通过光纤到达 OCS 的输入准直器 → 第一 MEMS 镜面 → 第二 MEMS 镜面 → 输出准直器 → 目标光纤——全程光学,无转发
  4. 光-电转换:在目标机架的光纤模块中,光信号转换回电信号
  5. ICI 链路层:电信号通过目标 Cube 的机架内部电缆和 PCB 走线,到达目标芯片 B 的 ICI 输入端口

除了两端的电-光转换模块,中间没有任何电设备参与。这个全光路径的延迟几乎完全由光纤长度决定(光速约 5 ns/m),加上 MEMS 镜面的微小光学路径长度差。


二、关键技术决策:设计选择的深层逻辑

TPU v4 的互联设计不是随意拍板的。每一个决策背后都有具体的工程约束和量化分析。这里我们展开 5 个关键技术决策。

2.1 决策一:为什么是 3D Torus 而不是 2D Torus 或 Dragonfly?

TPU v2/v3 用的是 2D Torus——每个芯片有 4 个邻居,对分带宽随芯片数 N 的平方根增长:O(√N)。当扩展到 1024 芯片以上时,2D Torus 的对分带宽趋于不足——尤其对 All-to-All 通信(嵌入层的典型通信模式)来说更是如此。

3D Torus 将这个比例提升到了 O(N^(2/3))——对分带宽在每个芯片数的比较基准上比 2D Torus 高出 2-4 倍(在同等芯片数下)。数值上,这意味着从 TPU v3(1024 芯片 2D Torus)升级到 TPU v4(4096 芯片 3D Torus)时,对分带宽的提升幅度远超简单的芯片数增长(4×)——实际上达到了 2-4× 的倍数增益。

为什么不是 Dragonfly 或 Clos?这两种拓扑在 HPC 超级计算机中很常见,但它们都需要专用的交换机层级——而这正是 Google 想要避免的。ICI 链路的优势在于它们是”胶水式”的直连——芯片到芯片,不需要经过外部交换机。这减少了延迟跳数、额外硬件成本和功耗。Google 对 3D Torus 的选择本质上是用更多维度替代更多交换机——维度多了,对分带宽就宽了,而交换机根本就不需要了。

至于为什么不直接到更高维度(4D 或更高),答案是物理限制。每个 TPU v4 有 6 条 ICI 链路——刚好对应 ±X、±Y、±Z 六个方向。更高维度意味着更多的芯片引脚、更大的封装面积、更高的 PCB 走线密度、以及更复杂的物理排布。在当前技术节点下,6 个链路是工程上合理的平衡点。

2.2 决策二:Cube 大小为什么是 4×4×4?

Cube 的大小选择——4×4×4(64 芯片)还是 8×8×8(512 芯片)——是一个空间-经济-可维护性的三角优化。

一个 4×4×4 Cube 连同其 16 个 CPU Host(每个 Host 带 4 个 TPU v4)刚好放入一个标准机架。这带来的好处是:机架内部的 ICI 链路全部使用被动电电缆——便宜、低功耗、无需光电转换。如果选择 8×8×8(512 芯片),则需要跨多个机架,相当比例的 ICI 链路将变为光链路——当链路长度超过电互联的物理极限时,必须转为光纤传输。而光模块(transceiver)的成本是电链路的 10 倍以上。

此外,43 Cube 的大小在调度上也更灵活。调度器不需要找到连续的 512 个空闲芯片——它可以组合 4 个来自不同物理位置的 43 Cube,通过 OCS 重新配置成一个逻辑上的 8×4×4 拓扑(128 芯片切片)。这种灵活性大大降低了碎片化,提高了集群利用率。

2.3 决策三:OCS 的成本-收益方程

OCS 的设计目标最初是解决可扩展性和可靠性的问题,但它在性能上的回报——尤其是对 LLM 训练——甚至超出了设计者的预期。

让我们仔细核算这笔账:

OCS 的成本端:48 个 Palomar OCS + 光模块 + 光纤基础设施,总成本 <5% 系统资本支出,总功耗 <3% 系统功耗。

OCS 的收益端是八项协同优势:

  1. 可扩展性(Scalability):4096 芯片,是 TPU v3(1024 芯片)的 4 倍。3D Torus + OCS 是这个规模级别可工程实现的唯一互联方案
  2. 可用性(Availability):OCS 可以像接线板一样绕过失效节点。当 CPU Host 可用性在 99.0%-99.5% 时(现实中正是这个水平),没有 OCS 的大型切片可用性极差;有了 OCS,各尺寸切片的可用性保持合理水平。系统整体达到 99.98% 可用性
  3. 灵活性(Modularity):支持 64 到 3072 芯片之间的多种拓扑形状,包括规则 3D Torus、扭曲 Torus 和”雪茄”形状(如 4×4×32 用于流水线并行)
  4. 性能(Performance):用户可以根据应用选择最优拓扑——在 ISCA 论文的表 3 中,仅靠改变拓扑和并行策略,LLM 训练吞吐量从 17.9 seq/s 提升到 41.3 seq/s(2.3×),GPT-3 预训练从 21.0 提升到 25.0(1.2×)
  5. 功耗(Power):MEMS 光学交换的功耗仅为电子交换的零头——OCS 是被动反射光而非主动处理数据包
  6. 调度简化(Scheduling):TPU v3 需要找到连续的物理芯片来组成切片;TPU v4 可以从超级计算机的任何位置挑选 Cube。切片大小甚至不需要是 2 的幂——可以是 4i×4j×4k,只要 0 < i <= j <= k 即可
  7. 部署速度(Deployment):Cube 是独立的部署单元,安装一个投入使用一个,不需等待全系统完成。装机时间大幅缩短,带来更好的投资回报
  8. 安全性(Security):OCS 可以在不同切片之间实现物理空气间隙(air-gap)隔离——不同的光路在物理上互不接触,使多租户共享成为可能

2.4 决策四:扭曲 Torus 的 All-to-All 增益

标准 3D Torus 有一个结构性问题:当切片尺寸不是完美立方体时(比如 4×4×8 而非 8×8×8),对分带宽不是最优的。

扭曲 Torus(Twisted Torus)通过重新布线部分 OCS 连接——实际上就是修改路由表——来减少最坏情况下的延迟并增加对分带宽。它不需要任何物理重新布线:因为 OCS 是软件可重配置的,在保持 Cube 内部的电连接不变的前提下,只改变 OCS 的光路径映射,就可以将 4×4×8 的规则矩形 Torus 变为扭曲 Torus。

TPU v4 中使用的扭曲配置来自 Camarero、Martinez 和 Beivide 的学术工作——k×k×2k 或 n×2n×2n 的配置。在生产中:

  • 4×4×8 切片:扭曲 Torus 比规则 Torus 的 All-to-All 吞吐量高 1.63×
  • 4×8×8 切片:扭曲 Torus 比规则 Torus 的 All-to-All 吞吐量高 1.31×

注意这个增益不只是针对嵌入层。在 LLM 训练中,许多集体通信操作(尤其是 AllReduce 的反向路径)也受对分带宽限制,扭曲 Torus 带来的对分带宽提升对它们同样有用。

在 Google 2022 年 11 月的一天抽样数据中(见 ISCA 论文表 2):43 Cube 以上的切片中,约 40% 选择了扭曲 Torus 拓扑。最流行的切片尺寸是 4×4×8_T(128 芯片,扭曲),占 16.0% 的使用量——高于 4×4×4(64 芯片常规 Torus)的 13.9%。这说明在许多生产场景下,扭曲 Torus 的实际价值得到了充分验证。

2.5 决策五:SparseCore — 嵌入加速与 MoE 通信的协同设计

TPU v4 中还有一个容易被忽略但对 MoE 意义深远的设计:SparseCore

SparseCore 是一个为嵌入(embedding)训练设计的专用数据流处理器,自 TPU v2 起就集成在芯片上。在 TPU v4 中,每个芯片有 4 个 SparseCore,它们以”芯海(sea-of-cores)”模型工作,通过 ICI 互联构成一个 128 TiB 的全局可寻址 HBM 空间。

SparseCore 和 MoE 有什么关系?

在推荐系统模型中,嵌入层产生的是 All-to-All 通信模式——每个 token 的嵌入向量需要在不同芯片间做分发和收集,通信的数据粒度细碎但总体量大。这和 MoE 的专家路由形成的 All-to-All 通信有着相同的底层特征:变量长度、稀疏、以对分带宽为瓶颈

TPU v4 通过将嵌入层放在 SparseCore 上运行,同时将密集计算(如 Transformer 层)放在 TensorCore 上运行,实现了两类操作的空间并行。这产生了一个对 MoE 有利的协同效应:SparseCore 已经建立了高效的 All-to-All 通信基础设施(包括 ICI 上的细粒度 scatter/gather 能力),当 LLM 中的 MoE 层产生类似的通信模式时,这个基础设施可以直接复用。

SparseCore 的设计本身也反映了以通信为中心的思想:

  • 16 个 Compute Tile(每 Tile 绑定一个 HBM 通道并支持多个并发内存访问)
  • 5 个 Cross-Channel Unit(在 16 个 Tile 的 Spmem 之间执行跨通道操作,如求和、排序、去重)
  • 专用的 Fetch Unit 和 Flush Unit(管理 HBM ↔ Spmem 之间的嵌入数据流)
  • 可编程的 8 宽 SIMD 向量处理单元(scVPU)

在性能实测中(TPU v4 vs TPU v3 on DLRM0, 128 chips),TPU v4 是 TPU v3 的 3.1×,是纯 CPU 配置的 30.1×。如果把嵌入数据放回 CPU 内存而非 SparseCore,TPU v4 的性能会下跌 5-7×——CPU 的内存带宽和 PCIe 成为更严重的瓶颈。这个 5-7× 的下跌非常精确地告诉了我们在 MoE 训练中为什么片上内存带宽和专用互联如此关键。


三、软件基础设施:让 4096 芯片可靠运行的协议栈

ISCA’23 的 TPU v4 论文主要聚焦硬件架构,但在同一个超级计算机上,2024 年 NSDI 上发表的另一篇论文 Resiliency at Scale: Managing Google’s TPUv4 Machine Learning Supercomputer 揭示了维持这个系统真正运行的软件全貌。

3.1 健康监控:healthd 的实时诊断

在每个 TPU v4 机器上都运行着一个名为 healthd 的后台守护进程,它持续监控:

  • 24 条单向 ICI 链路(接收和发送两个方向各 12 条 × 2)
  • TPU 与 CPU 之间的 PCIe 通道
  • 4 个 TPU ASIC 本身的健康状态(温度、功耗、时钟频率、错误寄存器)
  • 光模块的链路质量——通过测量误码率、光功率、波长偏移等参数

healthd 会将检测到的症状按严重程度排序。严重级别的症状会触发通知——通过 Borg 集群管理系统重新调度受影响的任务到健康的设备上。在每次用户任务开始前,还有两个预检(preflight check):

  • 端到端工作负载测试(实际跑一个微型计算来验证芯片和链路)
  • 基于意图的物理层指标验证(对比实测值和模型预期值,检测异常偏移)

TPU v4 集群的日均故障率:主机约 0.08%、ICI 线缆约 0.005%、OCS 约 0.04%。对于每个单独组件,这是非常低的数字——但在 4096 芯片的规模下,任何一天都有概率出现某种形式的硬件故障。关键不是消除故障,而是让故障不影响整体服务。这引出了下一个核心能力——容错路由。

3.2 libtpunet 协议栈:从物理层到事务层的全栈设计

libtpunet 是运行在每个用户任务内部的 ICI 网络设置和监控库。它的四层协议栈设计值得逐层分析:

物理层:libtpunet 启动时,首先扫描本地相邻 ICI 链路的连接状态——通过 BFS(广度优先搜索)自底向上发现整个拓扑。它会为每对芯片之间的 RDMA 通信分配唯一芯片 ID。物理通道建立后,libtpunet 持续监控链路状态并处理硬中断(MSI-X),一旦链路错误被检测到,libtpunet 可以触发在线链路排空(link drain)——在链路彻底断开前先将未完成传输的数据排空,然后触发任务重调度。这种”优雅降级”避免了训练任务的突然中断。

可靠数据层:在物理通道之上建立可靠传输——类似 RDMA 的语义。链路级别的流控缓冲区大小按照链路 RTT 成比例分配,确保在长距离光纤链路上(Cube 间走 OCS)不会因为缓冲区不足而带宽下降。时钟同步使用最小生成树算法,以最长的 ICI 路径 RTT 作为参考值来配置时钟偏移和同步参数。

路由层:这是整个协议栈中最精妙的部分。在无故障场景下,libtpunet 使用维序路由(Dimension-Ordered Routing, DOR)——数据包严格按照 X→Y→Z 的顺序沿维度前进。这保证只用两个虚拟通道就能实现无死锁。当包需要绕行半个维度(刚好在 Torus 环的对侧)时,路径选择根据源坐标的奇偶性判断——偶数走正方向,奇数走负方向——这样可以避免所有包拥堵同一方向。

当 OCS 失效或链路出现零星故障时,DOR 选择的路径可能不再可用。此时 libtpunet 切换到Wild First Routing(WFR)容错路由。WFR 的核心思想是:每个维度最多允许一次”野跳”(wild hop,即脱离标准维序的跳跃),之后恢复标准 DOR。为了防止死锁,WFR 要求野跳的维度顺序与标准 DOR 的维度顺序相反(”三明治规则”)。举个例子:在 X→Y→Z 的 DOR 下,使用 xyZYX 模式可以在避开 Z 方向故障的同时保持无死锁——先沿 X 和 Y 走到某个”绕行点”,再沿 Z 方向跳一次(野跳),然后沿 Y 和 X 退回正常路径方向。全部只用两个虚拟通道,维持了硬件上的 VC 数量限制。

事务层:虽然 ICI 的物理路径可能经过多跳,但给上层(TensorFlow/PyTorch/PATHWAYS 的计算图)展示的是一个全局可寻址的扁平地址空间——每个芯片只需要知道目标芯片 ID,而不需要知道中间经过哪些节点。ICI 交换器使用静态转发表(每对源-目的只有一条确定的路径),转发表的内容由离线 ILP(整数线性规划)求解器预计算。ILP 的目标是最大化所有到所有流量的吞吐量(即最大并发流问题),利用 Torus 的平移对称性来压缩变量空间——因为 Torus 是顶点对称的,一个源节点的最优路径集可以通过平移映射到所有源节点。预计算的转发表缓存在本地,libtpunet 在运行时加载对应的配置。

在 OCS 故障场景下,故障模式是 Cube 粒度的(一个 OCS 失效影响特定方向对的互联),ILP 求解仍然可以利用周期性将解空间限制在一组规范源节点中。实测结果是:单 OCS 失效时,规则 Torus 的 All-to-All 吞吐量降至约 93.4%(15/16 的理想到达率);而扭曲 Torus 因为通过路径截断(tiebreaking)拥有更大的路由灵活性,有时反而略微提升——更多可选路径意味着更好的负载均衡。

3.3 99.98% 可用性如何达成

在日均硬件故障率约 1% 的基础设施上,达成 99.98% 的系统可用性——也就是说一年的计划外停机时间不超过 1.75 小时——需要多层次的容错设计:

第一层:OCS 物理容错。OCS 的角色是”光学接线板”——如果某台主机出现故障,OCS 重新配置光路,将周围的健康 Cube 重新连接,逻辑上绕过故障节点。调度器重新打包切片,使用健康的 Cube 代替有故障的 Cube。

第二层:路由层容错(WFR)。当故障是零星的链路级别(而非整个 Cube 失效),WFR 允许数据包在避开故障链路的同时前往目的地。95% 的训练任务开启了容错 ICI 路由。

第三层:任务层容错。通过 healthd 的预检和监控,有问题的机器在任务分配前就被标识为不可用。如果任务运行中故障发生(虽然概率低但确实存在),libtpunet 的链路排空机制给 Borg 调度器留出了重调度任务的窗口。

第四层:分形调度。调度器不需要找到连续的物理芯片——只需要从整个超级计算机中挑选 n 个健康的 Cube,通过 OCS 在逻辑上组成所需的拓扑。这使得碎片化的可用资源也能被有效利用——即使 4096 芯片中有一部分不可用,调度器仍然能找到满足需求的集合。

3.4 OCS 投入生产:每天数万次重配置

NSDI 论文中一个被低估的数字:每台 TPU v4 Pod 的 Pod Manager 每天协调数万次 OCS 交叉连接(cross-connect)操作。这些操作包括:

  • 训练任务的拓扑建立(从默认网格到用户指定的 3D Torus 或扭曲 Torus 形状)
  • 任务结束后的拓扑清理
  • 故障恢复时的拓扑重配置
  • 跨租户安全隔离的物理连接调整

在 MS 级别完成一次 MEMS 微镜调整,每天数万次的频率——这需要高度的自动化编排。Pod Manager 管理着这一切:从接收调度器的拓扑需求、到将需求编译为每个 OCS 的镜面角度配置、到下发配置并验证链路质量、再到确认拓扑建立后通知 libtpunet 继续后续的协议栈层设置(时钟同步、路由表加载、会话建立)。整个流程从毫秒到秒级完成——远快于手动物理重新布缆(显然需要数小时到数天)。


四、横向对比:NVSwitch 的全交换 vs Torus+OCS 对 MoE 意味着什么

把 NVIDIA 的 NVSwitch 全互联方案和 Google 的 Torus+OCS 方案放在一起对比,不只是讨论”谁更快”——真正的问题是:这两种不同的互联哲学,在 MoE 这类对通信要求极高的负载面前,各自意味着什么?

4.1 对分带宽的两种思考

NVSwitch 的逻辑是”给所有人满带宽”——在 NVSwitch 域内(通常 8 卡,现在扩展至 72 卡),每对 GPU 之间的带宽是均匀的、满速的。All-to-All 在这个域内没有瓶颈:所有 GPU 可以同时以满带宽向所有其他 GPU 发送数据,NVSwitch 是一个非阻塞 crossbar。

Torus+OCS 用了另一种思路。Torus 的自然对分带宽是有限的——在 8×8×8(512 芯片)Torus 中,中间剖面上的链路数决定了理论峰值。但 OCS 允许拓扑重新配置——你可以将拓扑选择为”加重对分的方向”(比如用扭曲 Torus 增加对分链路的有效连接数)。更重要的是,你可以根据通信模式来选择拓扑:流水线并行更偏重小数据量、低延迟、”雪茄”形状(4×4×32);嵌入密集型的 All-to-All 更偏重大对分带宽的立方体形状(8×8×8);模型并行需要拓扑形状与并行参数映射对齐。NVSwitch 给了你一个”一刀切”的高带宽域,Torus+OCS 给了你一个”可以重新雕塑”的带宽域。

4.2 缩放路径的差异

NVSwitch 的缩放路径是增加交换机的层数。第一代 NVSwitch 用了 6 个交换机覆盖 8 个 GPU;第二代扩展到两个层次以容纳 256 个 GPU。但每增加一层交换,就增加了延迟和硬件成本。目前 NVSwitch 域内的 GPU 数量在增长(NVL72 含 72 GPU),但距离 4096 芯片的单域还有距离。

Torus+OCS 的缩放路径是增加 Cube 数量。从 1 个 Cube(64 芯片)到 64 个 Cube(4096 芯片),通过 OCS 做单一层级的拓扑映射——拓扑是一样的(只是更大),延迟的增长是对数级别(因为网络直径以 O(n^(1/3)) 增长)。成本的增长是线性的——多一个 Cube 就是多一个 Cube 的造价和配套的 48 根光纤连接。没有交换机层级的非线性膨胀。

不同之处在于:NVSwitch 的缩放保带宽但不保延迟(多级交换引入了额外的跳数);Torus+OCS 的缩放保延迟趋势但不保带宽(对分带宽随芯片数增长但每芯片分到的对分带宽逐渐下降)。

4.3 对 MoE All-to-All 的实际影响

在 MoE 的 All-to-All 通信中,最关键的瓶颈是对分带宽——能否在同一时刻让所有芯片以接近满速率向所有其他芯片发送 token 的 hidden states。NVSwitch 在它的域内解决了这个问题,但代价是域的大小有限、域外就回到 InfiniBand。Torus+OCS 在整个 4096 芯片域内都能做 All-to-All,但受限于 Torus 的结构性对分限制。

但是,OCS 在 MoE 场景下有一个 NVSwitch 不具备的优势:拓扑可以跟通信模式对齐。 如果训练的 MoE 模型同时使用数据并行和专家并行,那么 Torus 的三个维度可以被分别映射给这两种并行——一维用于数据并行的 AllReduce(对延迟敏感,需要邻近芯片之间的低延迟路径),另外两维用于专家并行的 All-to-All(对分带宽敏感,需要宽的双向切割面)。NVSwitch 虽然给了满速互联,但无法选择”哪个流量走哪条路”——它给的是一张对所有流量等价的单一大网。

这是两种哲学最深层的分歧。NVSwitch 相信”给够带宽,问题就会消失”——无论什么通信模式。Torus+OCS 相信”合理的带宽 + 灵活的分配 = 更好的整体效率”——带宽不是免费给的,而是在了解通信需求后定向分配。

PaLM 模型(540B 参数)在 50 天的 TPU v4 训练中,持续以约 57.8% 的峰值 FLOPS 运行——对于大模型训练来说,这是一个惊人的高效率。要知道,在 NVIDIA DGX SuperPOD 上,大模型训练达到 40-50% 的 MFU(Model FLOPS Utilization)已被认为是优秀的工程实践。TPU v4 和 OCS 的组合在不使用 Fat Tree 或 Clos 交换机网络的前提下,实现了顶尖水平的芯片利用率。


五、总结与清单:TPU v4 互联哲学的完整图景

核心教训

Google TPU v4 的互联设计挑战了一个被广泛接受的假设:高性能 AI 集群需要昂贵的、深层次的交换机网络。 它证明了一种替代方案——用更多的物理维度 + 光电路交换的可重配置拓扑,以更低的成本和功耗实现竞争力不亚于甚至超越同步竞品的训练效率。

这条路线有几个前提条件,不适用于所有组织:

  • 自研 ASIC 和互联:ICI 链路是 TPU 芯片上的原生接口,不需要网卡(GPU+IB 方案中 NIC 是额外的成本和功耗)。这需要从芯片层面设计互联。
  • 光学工程能力:Palomar OCS 是 Google 经过多年打磨的成果(源自数据中心网络 Jupiter 项目的光交换积累)。从 MEMS 制造到光模块成本控制,门槛非常高。
  • 全栈软件能力:libtpunet、healthd、ILP 求解器、Borg 调度器、PATHWAYS 编程框架——硬件只是故事的一半,让 OCS 和 Torus 在实际负载中体现价值的是这一整套软件栈。

但对于硬件架构师而言,即使不具备所有这些条件,TPU v4 的设计仍然传授了一种思维方式:与其试图消除拓扑的不均匀性,不如让拓扑适应负载的不均匀性。 这是一个可以作为未来互联设计指导原则的方向。

关键数字清单

指标 数值 意义
芯片数 4096 单个 TPU v4 超级计算机
Cube 大小 4×4×4(64 芯片) 电互联的基本构建块
Cube 数 64 通过 OCS 互联
OCS 数量 48 每个 Cube 6面 × 16 条链路 ÷ 2
OCS 端口数 136(128+8 备用) Palomar MEMS 微镜交换机
OCS 切换时间 毫秒级 适用于任务间拓扑重配置
OCS 成本占比 <5% 总系统成本 包括所有光学组件
OCS 功耗占比 <3% 总系统功耗 MEMS 被动维持功耗极低
ICI 单链路带宽 50 GB/s(400 Gbit/s) 每条链路,共 6 条/芯片
TPU v4 vs v3 性能 2.1× 同等芯片数
TPU v4 vs v3 能效 2.7× 性能/Watt
TPU v4 vs A100 性能 1.2×-1.7× 更快 同等系统规模
TPU v4 vs A100 功耗 1.3×-1.9× 更低 运行 MLPerf 基准
扭曲 Torus All-to-All 增益 1.63×(4×4×8),1.31×(4×8×8) 相对于规则 Torus
系统可用性 99.98% 硬件故障率 ~1%,容错机制保证
LLM 训练 MFU ~60%(PaLM 达到 57.8%) 50 天持续训练
SparseCore 面积/功耗 ~5% 芯片面积和功耗 嵌入加速 5-7×
SparseCore 全局 HBM 128 TiB 全系统 HBM 全局可寻址
容错路由采用率 95% 训练任务 通过 WFR 容忍链路故障
单 OCS 失效影响 All-to-All 降至 ~93.4% 规则 Torus,扭曲 Torus 更优

MoE 架构师的必问清单

如果你在设计一个 MoE 训练系统,面对互联的拓扑选择,以下问题值得在架构阶段逐个审视:

  1. All-to-All 的带宽需求怎样随专家数增长? 每个专家的容量因子 (capacity_factor) 是多少?如果 capacity_factor > 1,多余的带宽开销(padding)在你的互联上边际成本是多少?
  2. 你的系统同时用了几种并行方式? 如果数据并行和专家并行并存,能否像 Torus+OCS 那样将不同的维度(物理链路集合)分别分配给不同的通信模式?
  3. 你的互联拓扑能否在工作负载之间动态改变? 如果不同训练任务有不同的最优拓扑形状,当前的互联是固定的吗?重配置的代价是什么?
  4. 容错路由器对你的通信模式意味着什么? 当链路故障时,重路由后的有效对分带宽降低了多少?你的训练任务对这个降低有多敏感?
  5. 嵌入层和 MoE 专家层的 All-to-All 是否可以共享基础设施? TPU v4 的 SparseCore 证明了嵌入加速和 MoE 通信之间有天然的协同性——你的芯片上是否可以利用这种协同?
  6. 你的互联成本直接按端口数摊销还是按峰值带宽? OCS 的被动光交换模型使每端口的成本与带宽无关——更多带宽不需更贵的交换机。你的互联方案是否也具备这种带宽-成本脱钩特性?
  7. 系统部署和扩容的粒度是什么? 是必须整个集群一次性上线(像 TPU v3),还是可以逐个构建块逐步投入生产(像 TPU v4)?

在设计互联时,NVSwitch 追求的是”让拓扑对程序员不可见”——通信对任何人到任何人都是均匀的。TPU v4 走的是相反的方向——”让拓扑对程序员可见且可控制”——用选择拓扑的权力换取带宽成本的效率。两条路都是可行的。但当你开始直面 MoE 的 All-to-All 通信压力时——当通信开始吃掉绝大部分训练时间时——TPU v4 的路线提供了一种非常不同的思考框架:不要试图消除通信的不均匀性,而是找到一种方式,让系统的不均匀性匹配应用的不均匀性。 OCS 正是实现这种匹配所需的物理机制。