本文档介绍了 Google Kubernetes Engine (GKE) 中 Cloud TPU 的 动态切片。动态切片将 TPU 预配与工作负载 TPU 切片需求分离,从而提高 TPU 效率和灵活性。动态切片可将工作负载启动时间缩短多达 5 倍,将恢复时间缩短多达 4.5 倍,从而帮助提高 TPU 资源利用率。
动态切片适用于希望优化 TPU 利用率、缩短预配时间并提高大规模训练和推理工作负载的容错能力的机器学习 (ML) 工程师和平台工程师。
在阅读本文档之前,您应该熟悉以下内容:
- GKE 中的 TPU。
- TPU 全容量模式概览。动态切片功能仅适用于使用全容量模式的 TPU。
什么是动态切片?
动态切片是一项 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 倍。
- 加快作业启动速度:工作负载切片可以动态形成, 与创建静态节点池相比,作业启动延迟最多可缩短 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 自定义资源控制器,用于管理动态切片。切片 控制器管理 切片自定义 资源 的生命周期,该资源表示 动态切片(处理创建、持续监控和删除)。它还会将基础架构健康数据(主机、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 日之间,您必须手动 触发这些节点上的主机维护事件 ,然后才能使用子切片。
使用调度器进行动态切片
如需使用动态切片,您可以使用以下任一选项:
- 使用您自己的调度器来管理切片自定义资源。如果您有复杂的调度要求,或者想要将动态切片与现有调度基础架构集成,此选项非常有用。如需 开始使用,请参阅搭配自定义 调度器使用动态切片。
- 使用调度器自动创建 Slice 自定义资源。您可以 配置 Kueue 和 拓扑 感知调度 (TAS) 以自动创建 Slice 自定义资源。如需开始使用,请参阅 使用 Kueue 和 TAS 调度动态切片。
后续步骤
- 在全容量 模式下请求 TPU 容量。
- 使用 Kueue 和 TAS调度动态切片。
- 搭配自定义 调度器使用动态切片。