GKE 动态切片简介

本文档介绍了 Google Kubernetes Engine (GKE) 中适用于 Cloud TPU 的动态切片。动态切片通过将 TPU 预配与工作负载 TPU 切片需求解耦，提高了 TPU 效率和灵活性。动态切片可将工作负载启动时间缩短至原来的 1/5，并将恢复时间缩短至原来的 1/4.5，从而提高 TPU 资源利用率。

动态切片适用于希望优化 TPU 利用率、缩短预配时间并提高大规模训练和推理工作负载容错能力的机器学习 (ML) 工程师和平台工程师。

在阅读本文档之前，您应该先熟悉以下内容：

什么是动态切片？

动态切片是一项 GKE TPU 优化功能，可在调度工作负载时将物理和静态 TPU 硬件预配与运行时的 TPU 切片分配分离开来。

使用静态 TPU 拓扑时，GKE 节点池会绑定到您在创建期间配置的特定拓扑。由于此限制，如果静态池中的任何节点发生故障，整个节点池都会受到影响。您的工作负载拓扑还必须与节点池 TPU 拓扑完全匹配才能进行调度，这可能需要针对不同的工作负载频繁重新预配 TPU 基础设施。

动态切片使 GKE 能够在调度工作负载时在运行时配置 TPU 切片。您无需为每个新工作负载手动创建具有所需确切 TPU 拓扑的节点池，而是将整个 TPU 容量预配到精细的固定大小节点池中。对于 Ironwood (TPU7x)，每个节点池都包含一个 4x4x4 拓扑，也称为子块。子块表示由 16 个节点组成的 TPU 虚拟机组，没有活跃的芯片间互连 (ICI) 拓扑网格。在运行时，当工作负载被调度时，GKE 动态切片会重新配置 TPU 网络互连，以将多个节点池拼接成所需的大型 TPU 切片，或将节点池细分为较小的 TPU 切片。这种重新配置可在几秒钟内创建工作负载所需的精确多维拓扑。

动态切片的优势

在 GKE 中实现动态 TPU 架构具有以下优势：

  • 更快地恢复工作负载：如果发生物理硬件故障，GKE 会将问题隔离到单个分区。GKE 切片控制器会自动重塑活跃切片并重新配置网络，以将故障分区替换为正常运行的备用分区。与重新创建整个节点池相比，此过程可将弹性或平均恢复时间 (MTTR) 提高多达 4.5 倍。
  • 加速作业启动：工作负载切片可以动态形成，与创建静态节点池相比，作业启动延迟最多可缩短五倍。
  • 故障隔离：硬件故障会被隔离到发生故障的特定分区。这种隔离有助于保护集群中的其他并发作业免受级联故障的影响。
  • 优化资源利用率：动态切片有助于消除闲置或未充分利用的容量。由于切片配置为匹配工作负载需求，因此动态切片有助于最大限度地提高舰队利用率并最大限度地减少空闲硬件。
  • 声明式编排：动态切片使用 Kubernetes 原生自定义资源和注解，例如 JobSet 和 Kueue。此方法会自动管理底层网络和硬件编排。

动态切片配置

动态切片提供以下配置：

  • 动态超级切片：在工作负载调度时，将多个物理上独立、预配好的 TPU 节点池组合在一起，形成一个虚拟切片。例如，您可以将两个 4x4x4 节点池组合在一起，形成一个 4x4x8 拓扑。借助此功能，您可以创建等于或大于 4x4x4 拓扑的切片。此配置需要 1.35.2-gke.1842000 或更高版本。您可以使用动态超级切片来训练超出单个物理硬件块容量的大型基础模型。

    切片控制器负责协调光路交换 (OCS) 网络结构内的物理重新配置。通过动态重新配置 OCS，GKE 可跨独立硬件机架扩展芯片间互连 (ICI) 网络。从工作负载的角度来看，组合后的子块作为一个环形网格运行。在环形网格中，连接会环绕。右边缘的芯片直接连接回左边缘的芯片，顶部连接到底部。如果您将这种排列方式可视化，它会形成环面（环形）形状。

  • 动态子切片：在工作负载级别将单个预配的 TPU 节点池分区为多个较小的独立单元。借助动态子切片，您可以在单个子块中创建更小的拓扑，具体包括 2x2x1、2x2x2、2x2x4 和 2x4x4。例如，您可以将一个 4x4x4 子块划分为一个 2x2x4 子切片和两个 2x2x2 子切片。此配置需要快速渠道中的 1.36.0-gke.3712000 或更高版本。

    借助动态子切片，您可以在单个物理节点池运行多个工作负载。例如，您可以同时运行微调、实验和在线推理。子切片在这些子切片之间提供电气和网络隔离，这有助于隔离工作负载之间的故障或性能影响。

动态切片配置的主要特征

动态切片配置具有以下特征：

  • 节点池的增量预配：动态切片使用增量预配，这是一种容错的节点池预配模型。此模式会将您的所有 TPU 容量转换为由 16 节点 Ironwood (TPU7x) 虚拟机组构成的节点池，让您即使在部分立方体运行状况不佳的情况下，也能继续预配和利用资源。
  • 切片控制器：在 GKE 控制平面中运行的 Kubernetes 自定义资源控制器，用于管理动态切片。切片控制器负责管理 Slice 自定义资源的生命周期，该资源表示动态切片（处理创建、持续监控和删除）。它还会将基础设施健康数据（主机、ICI、OCS）传播到节点标签，以帮助识别健康状况良好的候选节点。
  • 切片自定义资源：表示逻辑切片，并启动节点间链路（ICI 和 OCS）的动态配置，以形成请求的 TPU 拓扑。此过程会将多个子块拼接在一起（超级切片），或在单个子块中隔离较小的拓扑（子切片）。您可以通过检查切片自定义资源的状态字段来检查动态切片形成的进度或健康状况。

要求

如需在 GKE 中使用动态切片，您必须满足以下要求：

  • 在快速渠道中使用以下任一版本的 Standard 集群：
    • 对于动态超级切片配置（拓扑等于或大于 4x4x4），请使用 1.35.2-gke.1842000 版或更高版本。
    • 对于动态子切片配置（小于 4x4x4 的拓扑），请使用 1.36.0-gke.3712000 版或更高版本。
  • 使用 Ironwood (TPU7x) 版本。
  • 为节点使用 Container-Optimized OS 操作系统映像。
  • 如需使用增量预配，请使用所有容量模式预留。全容量模式是 TPU Cluster Director 启用的一项功能。
  • 对于动态子切片，请确保您的节点有待处理的维护事件。监控实例是否有待处理的维护事件。如果您的任何节点有待处理的维护事件，且结束时间在 2026 年 9 月 18 日至 2026 年 9 月 30 日之间，则您必须先在这些节点上手动触发主机维护事件，然后才能使用子切片。

使用调度器进行动态切片

如需使用动态切片，您可以使用以下任一选项：

后续步骤