GKE 中的 TPU 简介

本页面介绍了 Cloud TPU 如何与 Google Kubernetes Engine (GKE) 搭配使用，包括术语、张量处理单元 (TPU) 的优势以及工作负载调度注意事项。TPU 是 Google 定制开发的应用专用集成电路 (ASIC)，用于加速使用 TensorFlow、PyTorch 和 JAX 等框架的机器学习工作负载。

本页面适用于平台管理员和运维人员，以及运行机器学习 (ML) 模型（具有大规模、长时间运行或由矩阵计算主导等特征）的数据和 AI 专家。如需详细了解我们在 Google Cloud内容中提及的常见角色和示例任务，请参阅常见的 GKE 用户角色和任务。

在阅读本页面之前，请确保您熟悉机器学习加速器的工作原理。如需了解详情，请参阅 Cloud TPU 简介。

使用 GKE 中 TPU 的优势

GKE 为 TPU 节点和节点池生命周期管理提供全面支持，包括创建、配置和删除 TPU 虚拟机。GKE 还支持 Spot 虚拟机和使用预留 Cloud TPU。如需了解详情，请参阅 Cloud TPU 使用情况选项。

使用 GKE 中的 TPU 的优势包括：

  • 一致的运行环境：您可以将单一平台用于所有机器学习和其他工作负载。
  • 自动升级：GKE 会自动更新版本，从而减少运营开销。
  • 负载均衡：GKE 会分发负载，从而缩短延迟时间并提高可靠性。
  • 响应式扩缩：GKE 会自动扩缩 TPU 资源，以满足工作负载的需求。
  • 资源管理：使用 Kueue（Kubernetes 原生作业排队系统），您可以使用排队、抢占、优先级和公平共享来管理组织内多个租户的资源。
  • 沙盒选项：GKE Sandbox 可帮助您使用 gVisor 保护工作负载。如需了解详情，请参阅 GKE Sandbox。

开始使用 Ironwood (TPU7x)

Ironwood (TPU7x) 是 Google 的第七代 TPU，专为大规模 AI 工作负载而设计。如需详细了解 Ironwood (TPU7x)，请参阅以下资源：

以下每种集群创建选项在集群配置和工作负载调度方面都提供了不同程度的便利性和灵活性：

  • 使用加速处理套件 (XPK) 快速创建 GKE 集群并运行工作负载，以进行概念验证和测试。如需了解详情，请参阅 XPK 文档。
  • 使用 Google Cloud CLI 手动创建 GKE 集群实例，以便精确自定义或扩展现有生产 GKE 环境

与 GKE 中的 TPU 相关的术语

本页面使用以下与 TPU 相关的术语：

  • Cloud TPU ICI 弹性：一项有助于提高光学链路和光学电路开关 (OCS) 的容错能力的功能，用于在立方体之间连接 TPU。如需了解详情，请参阅 TPU 架构。
  • TPU 立方体：一种由互连的 TPU 芯片组成的 4x4x4 拓扑。此参数仅适用于 3 元组 ({A}x{B}x{C}) 形式的拓扑。
  • TPU 类型：Cloud TPU 类型，例如 v5e。
  • TPU 切片：一组位于同一 TPU Pod 内的芯片，通过高速芯片间互连 (ICI) 进行连接。切片通过芯片或 TensorCore 来描述，具体取决于 TPU 版本。
  • TPU 切片节点：由具有一个或多个互连 TPU 芯片的单个虚拟机表示的 Kubernetes 节点。
  • TPU 切片节点池：集群内一组具有相同 TPU 配置的 Kubernetes 节点。
  • TPU 拓扑：TPU 切片中 TPU 芯片的数量和物理排列方式。
  • 原子性：GKE 将所有互连的节点视为一个单元。在扩缩操作期间，GKE 会将整个节点组缩减为 0 并创建新节点。如果节点组中的机器发生故障或终止，GKE 会作为一个新单元重新创建整个节点组。
  • 不可变：您无法手动向一组互连的节点添加新节点。不过，您可以创建具有所需 TPU 拓扑的新节点池，并在新节点池上调度工作负载。

TPU 切片节点池的类型

GKE 支持两种类型的 TPU 节点池：

TPU 类型和拓扑决定了 TPU 切片节点可以是多主机还是单主机。我们建议：

  • 对于大规模模型，请使用多主机 TPU 切片节点。
  • 对于小规模模型，请使用单主机 TPU 切片节点。
  • 对于大规模训练或推理，请使用 PathwaysPathways。Pathways 通过使单个 JAX 客户端能够跨多个大型 TPU 切片编排工作负载，简化了大规模机器学习计算。如需了解详情，请参阅 Pathways。

多主机 TPU 切片节点池

多主机 TPU 切片节点池是包含两个或更多互连 TPU 虚拟机的节点池。每个虚拟机都有一个与之连接的 TPU 设备。多主机 TPU 切片中的 TPU 通过高速互连 (ICI) 进行连接。创建多主机 TPU 切片节点池后，您无法向其添加节点。例如，您无法创建 v4-32 节点池，然后再向该节点池添加 Kubernetes 节点（TPU 虚拟机）。如需向 GKE 集群添加 TPU 切片，您必须创建新的节点池。

多主机 TPU 切片节点池中的虚拟机会被视为单个原子单元。如果 GKE 无法部署切片中的某个节点，则不会部署 TPU 切片节点中的任何节点。

如果多主机 TPU 切片中的节点需要维修，GKE 会关闭 TPU 切片中的所有虚拟机，从而强制逐出工作负载中的所有 Kubernetes Pod。当 TPU 切片中的所有虚拟机都已启动并运行后，Kubernetes Pod 便可以在新 TPU 切片中的虚拟机上进行调度。

下图显示了 v5litepod-16 (v5e) 多主机 TPU 切片。此 TPU 切片有四个虚拟机。TPU 切片中的每个虚拟机都有四个通过高速互连 (ICI) 连接的 TPU v5e 芯片，并且每个 TPU v5e 芯片都有一个 TensorCore：

多主机 TPU 切片图

下图显示了一个 GKE 集群，其中包含一个 TPU v5litepod-16 (v5e) TPU 切片（拓扑：4x4）和一个 TPU v5litepod-8 (v5e) 切片（拓扑：2x4）：

TPU v5e Pod 图

单主机 TPU 切片节点池

单主机切片节点池是包含一个或多个独立 TPU 虚拟机的节点池。每个虚拟机都有一个与之连接的 TPU 设备。虽然单主机切片节点池中的虚拟机可以通过数据中心网络 (DCN) 进行通信，但连接到虚拟机的 TPU 不会互连。

下图显示了一个包含 7 台 v4-8 机器的单主机 TPU 切片示例：

单主机切片节点池图

GKE 中的 TPU 特性

TPU 具有独特的特性，需要特殊规划和配置。

TPU 使用情况

为了在平衡工作负载性能的同时优化资源利用率和费用，GKE 支持以下 TPU 使用情况选项：

如需选择满足工作负载要求的消费选项，请参阅 GKE 中 AI/机器学习工作负载的加速器使用选项简介。

在 GKE 中使用 TPU 之前，请选择最符合工作负载要求的消费选项。

拓扑

拓扑定义了 TPU 切片内 TPU 的物理排列方式。GKE 会在二维或三维拓扑中配置 TPU 切片，具体取决于 TPU 版本。您可以采用每个维度中的 TPU 芯片数量的形式指定拓扑，如下所示：

对于在多主机 TPU 切片节点池中调度的 TPU v4、v5p 和 Ironwood (TPU7x)（预览版），您可以使用 3 元组 ({A}x{B}x{C}) 来定义拓扑，例如 4x4x4。{A}x{B}x{C} 的乘积定义了节点池中的 TPU 芯片数量。例如，您可以使用 2x2x2、2x2x4 或 2x4x4 等拓扑形式定义少于 64 个 TPU 芯片的小型拓扑。如果您使用的是具有超过 64 个 TPU 芯片的较大拓扑，则分配给 {A}、{B} 和 {C} 的值必须满足以下条件：

  • {A}、{B} 和 {C} 必须是四的倍数。
  • v4 支持的最大拓扑为 12x16x16，v5p 支持的最大拓扑为 16x16x24。
  • 分配的值必须保持 A ≤ B ≤ C 模式。例如 4x4x8 或 8x8x8。

机器类型

支持 TPU 资源的机器类型遵循以下命名惯例：其中包含 TPU 版本和每个节点切片的 TPU 芯片数量，例如 ct<version>-hightpu-<node-chip-count>t。例如，机器类型 tpu7x-standard-4t 支持 Ironwood (TPU7x)，且包含 4 个 TPU 芯片。

特权模式

如果您使用的是早于 1.28 版的 GKE 版本，则必须配置具有特殊功能的容器才能访问 TPU。在 Standard 模式集群中，您可以使用特权模式来授予此访问权限。特权模式会替换 securityContext 中的许多其他安全设置。如需了解详情，请参阅在不使用特权模式的情况下运行容器。

1.28 版及更高版本不需要特权模式或特殊功能。

GKE 中 TPU 的工作原理

Kubernetes 资源管理和优先级处理 TPU 上的虚拟机的方式与其他类型的虚拟机相同。如需请求 TPU 芯片，请使用资源名称 google.com/tpu：

resources:
  requests:
    google.com/tpu: 4
  limits:
    google.com/tpu: 4

使用 GKE 中的 TPU 时，请考虑以下 TPU 特性：

  • 一个虚拟机最多可以访问 8 个 TPU 芯片。
  • TPU 切片包含固定数量的 TPU 芯片，具体数量取决于您选择的 TPU 机器类型。
  • 请求的 google.com/tpu 数量必须等于 TPU 切片节点上的可用 TPU 芯片总数。GKE Pod 中请求 TPU 的任何容器都必须使用节点中的所有 TPU 芯片。否则，您的 Deployment 会失败，因为 GKE 无法使用部分 TPU 资源。以下面几种情况为例：
    • 具有 2x4 拓扑的机器类型 ct5lp-hightpu-4t 包含两个 TPU 切片节点，每个节点包含 4 个 TPU 芯片，共包含 8 个 TPU 芯片。使用此机器类型时，您：
    • 无法在此节点池中的节点上部署需要 8 个 TPU 芯片的 GKE Pod。
    • 可以部署两个各自需要 4 个 TPU 芯片的 Pod，每个 Pod 部署在此节点池中的两个节点之一上。
    • 采用拓扑 4x4 的 TPU v5e 在 4 个节点中有 16 个 TPU 芯片。选择此配置的 GKE Autopilot 工作负载必须在每个副本（1 到 4 个副本）中请求 4 个 TPU 芯片。
  • 在 Standard 集群中，一个虚拟机上可以调度多个 Kubernetes Pod，但每个 Pod 中只有一个容器可以访问 TPU 芯片。
  • 如需创建 kube-system Pod（例如 kube-dns），每个 Standard 集群必须至少有一个非 TPU 切片节点池。
  • 默认情况下，TPU 切片节点具有google.com/tpu 污点，以防止非 TPU 工作负载被调度到 TPU 切片节点上。 不使用 TPU 的工作负载会在非 TPU 节点上运行，从而为使用 TPU 的代码释放 TPU 切片节点上的计算资源。请注意，污点并不能保证 TPU 资源得到充分利用。
  • GKE 会收集 TPU 切片节点上运行的容器发出的日志。如需了解详情，请参阅日志记录。
  • Cloud Monitoring 提供 TPU 利用率指标（例如运行时性能）。如需了解详情，请参阅可观测性和指标。
  • 您可以使用 GKE Sandbox 对 TPU 工作负载进行沙盒处理。GKE Sandbox 适用于 TPU v4 及更高版本的模型。如需了解详情，请参阅 GKE Sandbox。

使用 TPU 自动创建节点池

节点池自动创建功能仅在特定 GKE 版本中支持以下 Cloud TPU：

  • TPU v3：1.31.0 及更高版本。
  • TPU v5 和 TPU v4：1.29.0 及更高版本。
  • TPU Trillium：1.31.1-gke.1146000 及更高版本。
  • Ironwood (TPU7x)（预览版）：1.34.1-gke.2541000 或更高版本。

所有 GKE 版本都支持其他 Cloud TPU 类型。

Cloud TPU 节点池自动扩缩

GKE 会以以下方式之一自动扩缩使用集群自动扩缩器自动创建或手动创建的 Cloud TPU 节点池：

  • 单主机 TPU 切片节点池：GKE 会在现有节点池中添加或移除 TPU 节点。节点池可以包含零到节点池最大大小（由 --max-nodes 和 --total-max-nodes 自动扩缩标志确定）之间的任意数量的 TPU 节点。节点池中的所有 TPU 节点都具有相同的机器类型和拓扑。如需详细了解如何创建单主机 TPU 切片节点池，请参阅创建单主机 TPU 切片节点池。
  • 多主机 TPU 切片节点池：GKE 以原子方式将节点池从零扩容到满足 TPU 拓扑所需的节点数。例如，对于机器类型为 ct5lp-hightpu-4t 且拓扑为 16x16 的 TPU 节点池，该节点池始终包含 64 个节点或零个节点。如果节点池中没有 TPU 工作负载，GKE 会缩减节点池。如需缩减节点池，GKE 会逐出所有调度的 Pod 并移除节点池中的所有节点。如需详细了解如何创建多主机 TPU 切片节点池，请参阅创建多主机 TPU 切片节点池。

工作负载政策

工作负载政策仅适用于 Ironwood (TPU7x)。

工作负载政策是一种资源政策，可用于配置底层基础设施的物理位置。工作负载政策提供分层配置选项，可实现更精细的控制。

Ironwood (TPU7x) 支持高吞吐量工作负载政策。此政策可让您执行以下操作：

  • 运行需要高吞吐量的分布式工作负载，例如高性能计算 (HPC) 或机器学习 (ML) 训练。底层基础架构配置为将 TPU 虚拟机并置，以缩短虚拟机之间的网络延迟。
  • 为 TPU 虚拟机定义维护策略，以最大限度地减少工作负载中断。

合集调度

只有 TPU Trillium 支持合集调度。

在 TPU Trillium 中，您可以使用合集调度对 TPU 切片节点进行分组。对这些 TPU 切片节点进行分组后，您可以更轻松地调整副本数量，以满足工作负载需求。 Google Cloud 会控制软件更新，以确保集合中的切片充足，始终可用于处理流量。

TPU Trillium 支持为运行推理工作负载的单主机和多主机节点池提供合集调度。以下内容介绍了合集调度行为如何取决于您使用的 TPU 切片类型：

  • 多主机 TPU 切片：GKE 会将多主机 TPU 切片分组以形成集合。每个 GKE 节点池都是此集合中的副本。如需定义集合，请创建多主机 TPU 切片并为集合分配唯一名称。如需向集合中添加更多 TPU 切片，请创建另一个具有相同集合名称和工作负载类型的多主机 TPU 切片节点池。
  • 单主机 TPU 切片：GKE 会将整个单主机 TPU 切片节点池视为一个集合。如需向集合中添加更多 TPU 切片，您可以调整单主机 TPU 切片节点池的大小。

合集调度具有以下限制：

  • 您只能为 TPU Trillium 调度集合。
  • 您只能在创建节点池时定义集合。
  • 不支持 Spot 虚拟机。
  • 包含多主机 TPU 切片节点池的集合必须为集合中的所有节点池使用相同的机器类型、拓扑和版本。

您可以在以下情况下配置合集调度：

后续步骤

如需了解如何设置 GKE 中的 Cloud TPU，请参阅以下页面：