*本文由张景涛翻译解读

谷歌在今年4月份的NSDI 2024上发表了《Resiliency at Scale:Managing Google’s TPUv4 Machine Learning Supercomputer》一文，详细介绍了谷歌的机器学习训练加速器TPUv4的设计和运营经验。（文末附论文下载）

随着机器学习（ML）模型的规模和复杂性不断增长，异构超级计算机的强大计算能力使得模型训练可以有效地沿批量、张量和流水线维度进行并行化。然而，ML超级计算机的硬件/软件生态系统在规模扩展时面临两个挑战：首先是有效地并行化模型训练工作负载，其次是维护计算资源的高可用性，从而为ML训练作业带来高吞吐量。谷歌TPUv4机器学习超级计算基础设施旨在应对这些挑战，包括多个硬件和软件组件，如TPU芯片、专有的芯片间互连（ICI）、光电路交换（OCSes）、Borg集群管理服务、Pod管理器、libtpunet软件库和healthd守护进程。

本文介绍了谷歌第三代机器学习训练加速器TPUv4（Tensor Processing Unit）的软件基础设施设计和运行经验。TPUv4以4096节点超级计算机的形式部署，采用定制的3D环状互连。文章重点介绍了TPUv4超级计算机在自动故障恢复和硬件恢复方面的功能，采用了软件定义网络（SDN）方法来管理高带宽的芯片间互连（ICI），并使用光电路交换动态配置路由，以绕过机器、芯片和链路故障。该基础设施能够检测故障并自动触发重新配置，最小化对运行工作负载的干扰，并启动受影响组件的修复工作流程。通过动态重新配置方法，TPUv4超级计算机实现了99.98%的系统可用性，处理了约1%的训练作业所经历的硬件故障。

可重配置机器学习超级计算机系统架构

从静态Pod架构中吸取的教训

在TPUv4 ML超级计算机之前，最先进的是TPUv2和TPUv3静态Pod——之所以称为“静态”，是因为它们具有不可重配置的固定ICI网格。TPUv2 Pod有256个TPU，连接在一个16x16的ICI环面上，而TPVv3有1024个TPU，连接在一个32x32环面上。还有一个TPUv3的扩展版本，将4个Pod组合成一个128x32的网格，具有有限的ICI路由能力。这种所谓的多Pod版本是与应用程序集合模式共同设计的，用于探索大型模型的扩展策略。

上图说明了在静态Pod中，随着模型大小的扩展，可用性面临的挑战。要训练一个模型，所有TPU进程必须同时处于活动状态，以便通过ICI集合操作同步更新它们的权重。任何一个失败或中断的进程都会中断整个训练过程。

为用户的作业寻找适当的计算资源面临以下挑战：

1、 硬件停机：定期对硬件、固件和软件在ICI链路、TPU芯片和CPU主机级别进行维护，可以将资源从可调度池中移除。对于一个拥有数千个TPU的超级计算机来说，影响任何一个组件的事件相对频繁，使得很难找到可用的资源集。此外，随着系统和应用程序在规模和复杂性的增加，意外故障发生的频率也更高。没有可重配置性，要获得需要1024个主机的作业的合理可用性，意味着每个单独的主机必须保持99.9%的可用性；引入可重配置的OCS将主机可用性要求降低到99%。

2、工作负载碎片整理：许多作业通常争夺Pod的可调度资源的不同子集。由于这些作业在不可预测的时间来来去去，有时Borg必须移动（抢占）较小的作业以为待处理的较大训练作业腾出连续的TPU。随着用户优先级的混合，调度复杂性加剧。有了基于OCS的可重配置性，Borg不需要太担心TPU资源的物理连续性。相反，任何一组空闲的立方体都可以通过OCS交叉连接，供用户作业使用。

3、部署前置时间：由于计算和网络资源的紧密耦合，静态Pod在所有硬件安装完成之前不可以使用。有了可重配置的Pod，一旦配置成功OCS，立方体就可以部署并立即使用。

上述挑战激发了对TPUv4 Pod架构的重新思考。

TPUv4: 基于OCS的可重配置性

TPUv4超级计算机采用了可重配置的系统架构，基于可编程的ICI协议和OCS及ICI交换，优化了系统的可用性和在规模化部署时的弹性。系统由多个cube组成，每个cube是一个硬件单元，包含64个TPU芯片，以4x4x4的3D网格排列。超级计算机或pod拥有64个cube，总计4096个TPU。

TPUv4的OCS可重配置性大大提高了可用性。即使它们不物理连续，训练作业也可以使用任何立方体，这减轻了竞争作业的资源碎片化。硬件故障将受影响的立方体从资源池中移除，但允许使用健康的立方体继续运行。选择16机故障容限粒度是为了平衡便利性（每个机架部署、电源和网络）的同时，保留一个小的故障影响范围。

可重配置性由相应的软件基础设施管理。每个作业启动都会引导软件根据所需的拓扑结构和立方体选择建立一个独特的OCS xconnect。芯片、链路和开关的大量数量还要求自动故障诊断、恢复、作业重新调度和容错ICI路由。使用OCS扩展TPUv4 Pod具有低成本：OCS和光纤的成本小于TPUv4 Pod总资本成本的5%，它们的运行功率小于Pod总功率的3%。TPUv4 OCS超级计算机的资本和运营成本明显低于使用InfiniBand等数据包交换机进行扩展的替代方案。

可编程ICI协议

TPUv4的ICI协议被设计为可编程的，以便软件可以处理可重配置性和弹性的操作复杂性。一个TPUv4 Pod是一个ICI域，其中任何两个TPU都可以相互RDMA。每个ICI链路可以承载50GBps的单向带宽。TPUv4采用3D ICI网络拓扑，以实现高双分割吞吐量、大系统规模和低延迟，同时保持低成本并支持通过集合操作并行化工作负载。

如上图所示，每个TPUv4芯片都有一定的计算能力、一些高带宽内存和一个实现各种ICI协议层的ICI开关。ICI协议促进了每个作业的网络分区，其中连接性、寻址、路由和流量控制都为每个作业设置，用户会话不会跨越作业边界。这样，每个作业就拥有它使用的所有链路的独占权，提高了安全性并消除了网络共享和拥塞控制的额外系统复杂性。

表1显示了协议层及其相应的软件代理。从底层向上，分别是物理层、可靠数据层、路由层和事务层。

TPUv4集群的自动化管理和运维

关键组件

• Cubes: 每个cube由16个TPU机器组成，每台机器组织4个TPU芯片，形成一个2x2x1的网格。
• Inter-Chip Interconnect (ICI): 高速互联，直接连接TPU以允许设备间的直接通信。
• Optical Circuit Switches (OCSes): 动态交叉连接不同cube的ICI，形成用户请求的环状拓扑结构。
• Borg: 集群管理服务，负责接受、调度和管理TPUv4作业。
• Pod Manager: 集群级软件服务，通过响应Borg调度决策来管理多cube的连接性。
• libtpunet：一个软件库，为每个TPUv4用户作业设置所请求的ICI网络拓扑。
• healthd：在pod中每个主机上运行的软件守护进程，持续监控机器硬件健康并向集群级软件系统报告。

核心流程

图5总结了TPUv4机器学习作业从启动到完成的整个生命周期。以下是该过程中的关键步骤：

1.用户作业请求：用户指定所需的3D拓扑结构，例如(4x, 4y, 4z)，以及其他元数据。
2.Borg集群调度器：接收所有作业请求并排队，等待资源分配。
3.资源分配：一旦作业有资格调度，Borg选择一组候选cube，然后发布xconnect请求。
4.Pod Manager：定期轮询Borg以了解任何待处理的xconnect请求。对于每个请求，它指导相关的OCS交换旋转它们的MEMS镜面以建立光ICI物理通道。
5.OCS xconnect完成：假设所有OCS xconnect正确完成，Pod Manager向Borg发送确认。
6.作业二进制分发：得到Pod Manager的批准后，Borg向选定的TPU机器集群分发作业二进制文件。
7.飞行前健康检查：首先运行健康检查以确保每个TPU机器的完整硬件健康（任何故障都会导致Borg重新调度到不同的cube）。
8.ICI网络设置：由libtpunet设置ICI网络（数据和路由层）。
9.XLA TPU编译器：采用由libtpunet构建的slice拓扑抽象，并为分布式训练生成自动并行化的TPU程序。
10.执行编译后的程序：在每台机器上，编译后的TPU二进制文件将通过PCIe发送到TPU，然后执行。
11.作业执行与监控：在训练期间，Fleet维护服务持续监控所有TPU机器的硬件和软件健康状况。任何异常情况都会触发通知到Borg，Borg再通知受影响的运行作业，以便它们可以（如果可能的话）写入最新的模型检查点。
12.故障检测与恢复：如果检测到故障，运行中的作业可以自动中断并重新调度。重新调度的作业将从最新的模型检查点恢复。
13.硬件维护：一旦作业完成或因故障重新调度，故障硬件将被识别并送入修复工作流程进行诊断和修复。

上述流程展示了TPUv4作业从提交到完成所经历的各个阶段，包括资源管理、网络配置、健康监测和故障恢复等关键环节。这一流程对于确保TPUv4超级计算机能够高效、可靠地运行机器学习训练作业至关重要。

通过上述设计和实现，TPUv4超级计算机能够自动适应硬件故障，实现高可用性，并为机器学习训练作业提供所需的计算资源。这些设计和实现的细节对于理解TPUv4超级计算机的高效运行至关重要。

验证结果

TPUv4超级计算机在生产中的操作经验表明，通过软件定义的ICI网络方法，实现了强大的机器和交换故障恢复能力。自2020年以来，该软件一直在生产中运行，为谷歌云集群和内部用户运行TPUv4超级计算机，并保持了99.98%的系统可用性。此外，通过光开关的使用，TPUv4超级计算机在成本上也具有优势，OCS和光纤成本不到TPUv4 pod总资本成本的5%，其运行功率不到pod总功率的3%。

下图展示了超级计算机硬件故障和恢复的周统计，包括TPU机器、ICI电缆和OCS。

结论

TPUv4超级计算机是一个E级4096芯片计算系统，解决了快速ML模型演化的可用性和可扩展性挑战。TPUv4提供了比前一代约2.1倍的性能，但也具有基于光电路交换的立方体重配置功能，并使用容错ICI路由以允许在交换机故障时运行。

本文描述了TPUv4的端到端软件基础设施，为拓扑、路由、调度、中断、监控和硬件健康管理提供了灵活性。TPUv4的软件定义ICI网络方法实现了对机器和交换故障的强大弹性。该软件自2020年以来一直在生产中运行，为Google Cloud集群和内部用户运行TPUv4超级计算机，并保持了99.98%的系统可用性。