在 GKE Standard 中部署 TPU 工作负载

标准

本页面介绍了如何在 Google Kubernetes Engine (GKE) 标准集群中使用 Cloud TPU 加速器 (TPU) 请求和部署大规模人工智能 (AI) 和机器学习 (ML) 模型训练、调优和推理工作负载。

在 GKE 中配置和部署 TPU 工作负载之前，您应该先熟悉以下概念：

须知事项

在开始之前，请确保您已执行以下任务：

启用 Google Kubernetes Engine API。

启用 Google Kubernetes Engine API

如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请运行 gcloud components update 以获取最新版本。
注意：对于现有 gcloud CLI 安装，请务必设置 compute/region 和 compute/zone 属性。通过设置默认位置，您可以避免 gcloud CLI 中出现以下错误：One of [--zone, --region] must be supplied: Please specify location。

GKE 中的 TPU 可用性

使用 GKE 创建和管理具有 TPU 的节点池。您可以使用这些专门构建的加速器来执行大规模 AI 模型训练、调优和推理。

请参阅 GKE 中支持的 TPU 版本列表。

规划 TPU 配置

根据您的机器学习模型及其所需的内存量规划 TPU 配置。以下是规划 TPU 配置时相关的步骤：

确保为按需或 Spot 虚拟机提供足够的配额

如果您要创建具有按需或 Spot 虚拟机的 TPU 切片节点池，则必须在您要使用的区域中拥有足够的 TPU 配额。

创建使用 TPU 预留的 TPU 切片节点池不需要任何 TPU 配额。¹ 对于预留的 TPU，您可以放心地跳过此部分。

在 GKE 中创建按需或 Spot TPU 切片节点池需要 Compute Engine API 配额。Compute Engine API 配额 (compute.googleapis.com) 与 Cloud TPU API 配额 (tpu.googleapis.com) 不同，使用 Cloud TPU API 创建 TPU 时需要后者。

如需查看 TPU 的 Compute Engine API 配额的限制和当前用量，请按照以下步骤操作：

转到 Google Cloud 控制台中的配额页面。

转到“配额”

在 过滤条件框中，执行以下操作：

选择服务属性，输入 Compute Engine API，然后按 Enter 键。
选择类型属性，然后选择配额。

选择名称属性，然后根据 TPU 版本和机器类型输入配额的名称。例如，如果您计划创建其机器类型以 ct5lp- 开头的按需 TPU v5e 节点，请输入 TPU v5 Lite PodSlice chips。

TPU 版本	机器类型开头为	按需实例的配额名称	Spot² 实例的配额名称
TPU v4	`ct4p-`	`TPU v4 PodSlice chips`	`Preemptible TPU v4 PodSlice chips`
TPU v5e	`ct5l-`	`TPU v5 Lite Device chips`	`Preemptible TPU v5 Lite Device chips`
TPU v5e	`ct5lp-`	`TPU v5 Lite PodSlice chips`	`Preemptible TPU v5 Lite PodSlice chips`
TPU v5p	`ct5p-`	`TPU v5p chips`	`Preemptible TPU v5p chips`

选择维度（例如位置）属性，然后输入 region:，后跟您计划在其中创建 GKE 中 TPU 的区域的名称。例如，如果您计划在可用区 us-west4-a 中创建 TPU 切片节点，请输入 region:us-west4。TPU 配额是区域性的，因此同一区域内的所有可用区都使用同一 TPU 配额。

如果没有配额与您输入的过滤条件匹配，则说明项目尚未获得所需区域的任何指定配额，您必须申请增加 TPU 配额。

创建 TPU 切片节点池时，请使用 --reservation 和 --reservation-affinity=specific 标志创建预留实例。购买承诺时，TPU 预留可供使用。↩
创建 TPU 切片节点池时，请使用 --spot 标志来创建 Spot 实例。 ↩

确保预留可用性

创建预留 TPU 切片节点池（使用预留的 TPU 切片节点池）不需要任何 TPU 配额。但是，在创建节点池时，预留必须具有足够的可用或未使用的 TPU 芯片。

如需查看项目中存在哪些预留，请查看预留列表。

如需查看 TPU 预留中可用的 TPU 芯片数，请查看预留的详细信息。

创建集群

在具有可用 TPU 的区域中创建 Standard 模式的 GKE 集群。我们建议您使用区域级集群，可提供 Kubernetes 控制平面的高可用性。您可以使用 Google Cloud CLI 或 Google Cloud 控制台。

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

请替换以下内容：

CLUSTER_NAME：新集群的名称。
LOCATION：具有可用 TPU 容量的区域。
VERSION：GKE 版本，必须支持您要使用的机器类型。请注意，默认 GKE 版本可能无法为您的目标 TPU 提供可用性。如需了解 TPU 机器类型可用的最低 GKE 版本，请参阅 GKE 中的 TPU 可用性

创建节点池

单主机 TPU 切片

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台创建单主机 TPU 切片节点池。

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE

请替换以下内容：

POOL_NAME：新节点池的名称。
LOCATION：基于您要使用的 TPU 版本的可用区名称：
- 对于 TPU v4，请使用 us-central2-b。
- 对于以 ct5l- 开头的 TPU v5e 机器类型，请使用 us-central1-a 或 europe-west4-b。
- 对于以 ct5lp- 开头的 TPU v5e 机器类型，请使用 us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b 或 europe-west4-a。
- 对于 TPU v5p，请使用 us-east1-d、us-east5-a 或 us-east5-c。
如需了解详情，请参阅 GKE 中的 TPU 可用性。
CLUSTER_NAME：集群的名称。
NODE_ZONE：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。
MACHINE_TYPE：用于节点的机器类型。如需详细了解与 TPU 兼容的机器类型，请使用 TPU 配置映射中的表。

（可选）您还可以使用以下标志：

--num-nodes=NUM_NODES：每个可用区的节点池中的初始节点数。如果您省略此标志，则默认值为 3；如果使用 --enable-autoscaling 标志为节点池启用了自动扩缩功能，我们建议您将 NUM_NODES 设置为 0，因为自动扩缩器在工作负载需要时会预配更多节点。
--reservation=RESERVATION_NAME：GKE 在创建节点池时使用的预留的名称。如果您省略此标志，则 GKE 会使用可用的 TPU。如需详细了解 TPU 预留，请参阅 TPU 预留。
--enable-autoscaling：创建启用了自动扩缩功能的节点池。需要使用以下其他标志：
- --total-min-nodes=TOTAL_MIN_NODES：节点池中的所有节点数下限。
- --total-max-nodes=TOTAL_MAX_NODES：节点池中的所有节点数上限。
- --location-policy=ANY：优先使用未使用的预留，并降低 Spot 虚拟机的抢占风险。
--spot：设置节点池以对节点池中的节点使用 Spot 虚拟机。创建节点池后，此设置将无法更改。

如需查看您可以指定的所有标志的完整列表，请参阅 gcloud container clusters create 参考文档。

Terraform

确保您使用 google 提供程序 4.84.0 版或更高版本。
将以下块添加到 Terraform 配置中：

resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }
}

请替换以下内容：

NODE_POOL_RESOURCE_NAME：Terraform 模板中的节点池资源的名称。
PROJECT_ID：您的项目 ID。
CLUSTER_NAME：现有集群的名称。
POOL_NAME：要创建的节点池的名称。
CLUSTER_LOCATION：集群的计算可用区。指定 TPU 版本可用的区域。如需了解详情，请参阅选择 TPU 版本和拓扑。
NODE_ZONES：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。
MACHINE_TYPE：要使用的 TPU 机器的类型。如需查看与 TPU 兼容的机器类型，请使用 TPU 配置映射中的表格。

（可选）您还可以使用以下变量：

autoscaling：创建启用了自动扩缩功能的节点池。对于单主机 TPU 分片，GKE 会在 TOTAL_MIN_NODES 和 TOTAL_MAX_NODES 值之间进行扩缩。
- TOTAL_MIN_NODES：节点池中的所有节点数下限。除非同时指定了自动扩缩，否则此字段是可选字段。
- TOTAL_MAX_NODES：节点池中的所有节点数上限。除非同时指定了自动扩缩，否则此字段是可选字段。
RESERVATION_NAME：如果您使用 TPU 预留，则这是创建节点池时使用的预留资源的标签列表。如需详细了解如何填充 reservation_affinity 字段中的 RESERVATION_LABEL_VALUES，请参阅 Terraform 提供程序。
spot：设置节点池以对 TPU 节点使用 Spot 虚拟机。创建节点池后，此设置将无法更改。如需了解详情，请参阅 Spot 虚拟机。

控制台

如需创建具有 TPU 的节点池，请执行以下操作：

转到 Google Cloud 控制台中的 Google Kubernetes Engine 页面。

前往 Google Kubernetes Engine
在集群列表中，点击您要修改的集群的名称。
点击 添加节点池。
在节点池详情部分中，勾选指定节点位置复选框。
根据您要使用的 TPU 版本选择可用区：
- 对于 TPU v4，请使用 us-central2-b。
- 对于以 ct5l- 开头的 TPU v5e 机器类型，请使用 us-central1-a 或 europe-west4-b。
- 对于以 ct5lp- 开头的 TPU v5e 机器类型，请使用 us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b 或 europe-west4-a。
- 对于 TPU v5p，请使用 us-east1-d、us-east5-a 或 us-east5-c。
在导航窗格中，点击节点。
在机器配置部分中，选择 TPU。
在系列下拉菜单中，选择以下选项之一：
- CT4P：TPU v4
- CT5LP：TPU v5e
- CT5P：TPU v5p
在机器类型下拉菜单中，选择要用于节点的机器的名称。使用 TPU 配置映射表可了解如何定义用于创建单主机 TPU 切片节点池的机器类型和 TPU 拓扑。
在 TPU 拓扑下拉菜单中，选择 TPU 切片的物理拓扑。
在需要更改对话框中，点击进行更改。
确保启动磁盘类型为标准永久性磁盘或 SSD 永久性磁盘。
（可选）选中在 Spot 虚拟机上启用节点复选框，以对节点池中的节点使用 Spot 虚拟机。
点击创建。

多主机 TPU 切片

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台创建多主机 TPU 切片节点池。

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

请替换以下内容：

POOL_NAME：新节点池的名称。
LOCATION：基于您要使用的 TPU 版本的可用区名称：
- 对于 TPU v4，请使用 us-central2-b。
- 以 ct5l- 开头的 TPU v5e 机器类型绝不会是多主机。
- 对于以 ct5lp- 开头的 TPU v5e 机器类型，请使用 us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b 或 europe-west4-a。
- 对于以 ct5p- 开头的 TPU v5p 机器类型，请使用 us-east1-d、us-east5-a 或 us-east5-c。
如需了解详情，请参阅 GKE 中的 TPU 可用性。
CLUSTER_NAME：集群的名称。
NODE_ZONE：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。
MACHINE_TYPE：用于节点的机器类型。如需详细了解可用机器类型，请参阅 TPU 配置映射。
TPU_TOPOLOGY：TPU 切片的物理拓扑。拓扑格式取决于 TPU 版本，如下所示：
- TPU v4 或 v5p：使用 3 元组 ({A}x{B}x{C}) 定义拓扑，例如 4x4x4。
- TPU v5e：使用 2 元组 ({A}x{B}) 定义拓扑，例如 2x2。
如需了解详情，请参阅拓扑。
NUM_NODES：节点池中的节点数。该值必须为零或 TPU_TOPOLOGY ({A}x{B}x{C}) 中定义的值的乘积除以每个虚拟机中的芯片数量。对于多主机 TPU v4 和 TPU v5e，每个虚拟机中的芯片数量为 4。因此，如果 TPU_TOPOLOGY 为 2x4x4（每个虚拟机中有四个芯片的 TPU v4），则 NUM_NODES 为 32/4，等于 8。

（可选）您还可以使用以下标志：

RESERVATION_NAME：GKE 在创建节点池时使用的预留的名称。如果您省略此标志，则 GKE 会使用可用的 TPU 切片节点池。如需详细了解 TPU 预留，请参阅 TPU 预留。
--spot：设置节点池以对 TPU 切片节点使用 Spot 虚拟机。创建节点池后，此设置将无法更改。如需了解详情，请参阅 Spot 虚拟机。
--enable-autoscaling：创建启用了自动扩缩功能的节点池。当 GKE 扩缩多主机 TPU 切片节点池时，它会以原子方式将节点池从零扩容到大小上限。
- MAX_NODES：节点池的大小上限。如果提供了 --enable-autoscaling，则必须使用 --max-nodes 标志，且值必须等于 TPU_TOPOLOGY 中定义的值的乘积 ({A}x{B}x{C}) 除以每个虚拟机中的芯片数量。

Terraform

确保您使用 google 提供程序 4.84.0 版或更高版本。
将以下块添加到 Terraform 配置中：
```
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}
```
请替换以下内容：
- NODE_POOL_RESOURCE_NAME：Terraform 模板中的节点池资源的名称。
- PROJECT_ID：您的项目 ID。
- CLUSTER_NAME：要在其中添加节点池的现有集群的名称。
- POOL_NAME：要创建的节点池的名称。
- CLUSTER_LOCATION：集群的计算位置。我们建议您使用区域级集群，以提高 Kubernetes 控制平面的可靠性。您还可以使用可用区级集群。如需了解详情，请参阅选择 TPU 版本和拓扑。
- NODE_ZONES：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。
- NUM_NODES：节点池中的节点数。值必须为零或 TPU 芯片数量的乘积除以 4，因为在多主机 TPU 切片中，每个 TPU 切片节点都有 4 个芯片。例如，如果 TPU_TOPOLOGY 为 4x8，则有 32 个芯片，这意味着 NUM_NODES 必须为 8。如需详细了解 TPU 拓扑，请使用 TPU 配置映射中的表格。
- TPU_TOPOLOGY：指示所需的 TPU 切片物理拓扑。拓扑格式取决于您使用的 TPU 版本：
  - 对于 TPU v4：使用 3 元组 ({A}x{B}x{C}) 定义拓扑，例如 4x4x4。
  - 对于 TPU v5e：使用 2 元组 ({A}x{B}) 定义拓扑，例如 2x2。
（可选）您还可以使用以下变量：
- RESERVATION_NAME：如果您使用 TPU 预留，则这是创建节点池时使用的预留资源的标签列表。如需详细了解如何填充 reservation_affinity 字段中的 RESERVATION_LABEL_VALUES，请参阅 Terraform 提供程序。
- autoscaling：创建启用了自动扩缩功能的节点池。当 GKE 扩缩多主机 TPU 切片节点池时，它会以原子方式将节点池从零扩容到大小上限。
  - MAX_NODES：节点池的大小上限。该值必须等于 TPU_TOPOLOGY ({A}x{B}x{C}) 中定义的值的乘积除以每个虚拟机中的芯片数量。
- spot：可让节点池对 TPU 切片节点使用 Spot 虚拟机。创建节点池后，此设置将无法更改。如需了解详情，请参阅 Spot 虚拟机。

控制台

如需创建具有 TPU 的节点池，请执行以下操作：

转到 Google Cloud 控制台中的 Google Kubernetes Engine 页面。

前往 Google Kubernetes Engine
在集群列表中，点击您要修改的集群的名称。
点击 添加节点池。
在节点池详情部分中，勾选指定节点位置复选框。
根据您要使用的 TPU 版本选择可用区：
- 对于 TPU v4，请使用 us-central2-b。
- 以 ct5l- 开头的 TPU v5e 机器类型绝不会是多主机。
- 对于以 ct5lp- 开头的 TPU v5e 机器类型，请使用 us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b 或 europe-west4-a。
- 对于以 ct5p- 开头的 TPU v5p 机器类型，请使用 us-east1-d、us-east5-a 或 us-east5-c。
在导航窗格中，点击节点。
在机器配置部分中，选择 TPU。
在系列下拉菜单中，选择以下选项之一：
- CT4P：适用于 TPU v4。
- CT5LP：适用于 TPU v5e。
在机器类型下拉菜单中，选择要用于节点的机器的名称。使用 TPU 配置映射表可了解如何定义用于创建多主机 TPU 切片节点池的机器类型和 TPU 拓扑。
在 TPU 拓扑下拉菜单中，选择 TPU 切片的物理拓扑。
在需要更改对话框中，点击进行更改。
确保启动磁盘类型为标准永久性磁盘或 SSD 永久性磁盘。
（可选）选中在 Spot 虚拟机上启用节点复选框，以对节点池中的节点使用 Spot 虚拟机。
点击创建。

预配状态

如果 GKE 由于可用 TPU 容量不足而无法创建 TPU 切片节点池，GKE 会返回一条错误消息，指示由于容量不足而无法创建 TPU 切片节点。

如果您要创建单主机 TPU 切片节点池，则错误消息将类似于以下内容：

2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.

如果您要创建多主机 TPU 切片节点池，则错误消息将类似于以下内容：

The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.

您的 TPU 预配请求可能会在队列中长时间保留，在队列中时将一直保持“正在预配”状态。

具备可用容量后，GKE 会创建尚未创建的其余节点。

如果您需要更快地需要容量，请考虑使用 Spot 虚拟机，但请注意，Spot 虚拟机消耗的配额与按需实例不同。

您可以通过删除 TPU 切片节点池来删除已加入队列的 TPU 请求。

在 TPU 切片节点上运行工作负载

工作负载准备

TPU 工作负载具有以下准备要求。

JAX、PyTorch 和 TensorFlow 等框架使用 libtpu 共享库访问 TPU 虚拟机。libtpu 包括 XLA 编译器、TPU 运行时软件和 TPU 驱动程序。每个 PyTorch 和 JAX 版本都需要特定的 libtpu.so 版本。如需使用 GKE 中的 TPU，请确保您使用以下版本：

TPU 类型	`libtpu.so`版本
TPU v5e `tpu-v5-lite-podslice` `tpu-v5-lite-device`	推荐的 jax[tpu] 版本：v0.4.9 或更高版本推荐的 torchxla[tpuvm] 版本：v2.1.0 或更高版本。
TPU v5p `tpu-v5p-slice`	推荐的 jax[tpu] 版本：0.4.19 或更高版本。推荐的 torchxla[tpuvm] 版本：建议使用 2023 年 10 月 23 日的夜间版本 build。
TPU v4 `tpu-v4-podslice`	jax[tpu]：v0.4.4 或更高版本 torchxla[tpuvm]：v2.0.0 或更高版本

为请求 TPU 资源的容器设置以下环境变量：
- TPU_WORKER_ID：每个 Pod 的唯一整数。此 ID 表示 TPU 切片中的唯一工作器 ID。此字段支持的值范围从 0 到 Pod 数量减去 1。
- TPU_WORKER_HOSTNAMES：需要在切片内相互通信的 TPU 虚拟机主机名或 IP 地址的英文逗号分隔列表。切片中的每个 TPU 虚拟机都应该有一个对应的主机名或 IP 地址。IP 地址或主机名列表按 TPU_WORKER_ID 排序且为零索引。
- 在工作负载清单中，添加 Kubernetes 节点选择器，以确保 GKE 根据您定义的 TPU 机器类型和 TPU 拓扑调度 TPU 工作负载：
```
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
  
```
  替换以下内容：
  - TPU_ACCELERATOR：TPU 加速器的名称：
    - 对于 TPU v4，请使用 tpu-v4-podslice
    - 对于以 ct5l- 开头的 TPU v5e 机器类型，请使用 tpu-v5-lite-device
    - 对于以 ct5lp- 开头的 TPU v5e 机器类型，请使用 tpu-v5-lite-podslice
    - 对于 TPU v5p，请使用 tpu-v5p-slice
  - TPU_TOPOLOGY：TPU 切片的物理拓扑。拓扑格式取决于 TPU 版本，如下所示：
    - TPU v4：使用 3 元组 ({A}x{B}x{C}) 定义拓扑，例如 4x4x4。
    - TPU v5e：使用 2 元组 ({A}x{B}) 定义拓扑，例如 2x2。
    - TPU v5p：使用 3 元组 ({A}x{B}x{C}) 定义拓扑，例如 4x4x4。

完成工作负载准备后，您可以运行使用 TPU 的 Job。

以下部分举例说明了如何运行使用 TPU 执行简单计算的 Job。

示例 1：运行显示 TPU 切片节点池中可用 TPU 芯片数量的工作负载

以下工作负载会返回多主机 TPU 切片中所有节点的 TPU 芯片数量。如需创建多主机切片，工作负载需要具有以下参数：

TPU 版本：TPU v4
拓扑：2x2x4

选择此版本和拓扑会生成多主机切片。

将以下清单保存为 available-chips-multihost.yaml：

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4

部署清单：
```
kubectl create -f available-chips-multihost.yaml
```
GKE 会运行具有四个虚拟机的 TPU v4 切片（多主机 TPU 切片）。该切片具有 16 个互连 TPU 芯片。

验证作业是否创建了四个 Pod：

kubectl get pods

输出类似于以下内容：

NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s

获取其中一个 Pod 的日志：
```
kubectl logs POD_NAME
```
将 POD_NAME 替换为创建的某个 Pod 的名称。例如 tpu-job-podslice-0-5cd8r。

输出类似于以下内容：
```
TPU cores: 16
```

示例 2：运行显示 TPU 切片中可用 TPU 芯片数量的工作负载

以下工作负载是一个静态 Pod，显示了连接到特定节点的 TPU 芯片的数量。如需创建单主机节点，工作负载需要具有以下参数：

TPU 版本：TPU v5e
拓扑：2x4

选择此版本和拓扑会生成单主机切片。

将以下清单保存为 available-chips-singlehost.yaml：

apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    securityContext:
      privileged: true
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8

部署清单：
```
kubectl create -f available-chips-singlehost.yaml
```
GKE 会预配具有八个单主机 TPU 切片的节点，这些切片使用 TPU v5e。每个 TPU 节点都有八个 TPU 芯片（单主机 TPU 切片）。
获取相应 Pod 的日志：
```
kubectl logs tpu-job-jax-v5
```
输出类似于以下内容：
```
Total TPU chips: 8
```

使用加速器（GPU 和 TPU）升级节点池

GKE 会自动升级 Standard 集群，包括节点池。如果您希望节点更快地升级到更高版本，也可以手动升级节点池。如需控制集群升级的方式，请使用发布渠道、维护窗口和排除项以及发布顺序。

您还可以为节点池配置节点升级策略，例如超额配置升级或蓝绿升级。通过配置这些策略，您可以确保节点池的升级方式能够在速度和环境中断之间实现最佳平衡。对于多主机 TPU 切片节点池，GKE 会在单一步骤中以原子方式重新创建整个节点池，而不是使用配置的节点升级策略。如需了解详情，请参阅 GKE 中的 TPU 相关术语中的“原子性”定义。

使用节点升级策略会暂时要求 GKE 预配更多资源，具体取决于配置。如果 Google Cloud 对节点池的资源具有容量限制（例如，您尝试使用 GPU 或 TPU 创建更多节点时出现资源可用性错误），请参阅在资源受限的环境中升级。

清理

为避免因本指南中使用的资源而向您的 Google Cloud 账号收取费用，请考虑删除不再有调度工作负载的 TPU 切片节点池。如果必须正常终止工作负载运行，请使用 kubectl drain 清理工作负载，然后再删除节点。

删除 TPU 分片节点池：
```
gcloud container node-pools delete POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME
```
替换以下内容：
- POOL_NAME：节点池的名称。
- CLUSTER_NAME：集群的名称。
- LOCATION：集群的计算位置。

其他配置

以下各部分介绍了可应用于 TPU 工作负载的其他配置。

多切片

您可以在多切片中汇总较小的时间片，以处理较大训练工作负载。如需了解详情，请参阅 GKE 中的多切片 TPU。

迁移 TPU 预留

如果您有现有的 TPU 预留，则必须先将 TPU 预留迁移到新的基于 Compute Engine 的预留系统。您还可以创建基于 Compute Engine 的预留系统，而无需迁移。如需了解如何迁移 TPU 预留，请参阅 TPU 预留。

日志记录

GKE 节点（包括 TPU 虚拟机）上运行的容器发出的日志会由 GKE 日志记录代理收集，发送到 Logging，并且显示在 Logging 中。

使用 GKE 节点自动预配功能

您可以将 GKE 配置为自动创建和删除节点池，以满足 TPU 工作负载的资源需求。如需了解详情，请参阅配置 Cloud TPU。

TPU 切片节点自动修复

如果多主机 TPU 切片节点池中的 TPU 切片节点健康状况不佳，则系统会重新创建整个节点池。而在单主机 TPU 切片节点池中，系统只会自动修复健康状况不佳的 TPU 节点。

导致 TPU 切片节点健康状况不佳的条件包括：

任何具有通用节点条件的 TPU 切片节点。
任何不可分配 TPU 数量大于零的 TPU 切片节点。
TPU 切片中的任何已停止（由于抢占）或终止的虚拟机实例。
节点维护：如果多主机 TPU 切片节点池中的任何 TPU 切片节点因主机维护而关闭，GKE 会重新创建整个 TPU 切片节点池。

您可以在操作历史记录中查看修复状态（包括失败原因）。如果失败是由配额不足引起的，请与您的 Google Cloud 客户代表联系，以增加相应的配额。

配置 TPU 分片节点的正常终止

在控制平面运行 1.29.1-gke.1425000 或更高版本的 GKE 集群中，TPU 切片节点支持 SIGTERM 信号，用于提醒节点即将关停。在 TPU 节点中，即将关停通知的可配置时间最长为五分钟。

如需将 GKE 配置为在此通知时间范围内正常终止工作负载，请按照管理 GKE 节点中断情况（适用于 GPU 和 TPU）中的步骤操作。

在不使用特权模式的情况下运行容器

如果您的 TPU 切片节点运行的版本低于 1.28，请参阅以下部分：

在 TPU 切片中的虚拟机上运行的容器需要访问更高的锁定内存限制，以便驱动程序可以通过直接内存访问 (DMA) 与 TPU 芯片进行通信。如需实现这一目标，您必须配置更高的 ulimit。如果您想缩小容器的权限范围，请完成以下步骤：

修改 securityContext 以包含以下字段：

securityContext:
  capabilities:
    add: ["SYS_RESOURCE"]

在将工作负载设置为使用 TPU 资源之前，请在容器内运行以下命令以增加 ulimit：
```
ulimit -l 68719476736
```

注意：对于 TPU v5e，运行 1.27.4-gke.900 版及更高版本的集群可以在无特权模式下运行容器。

可观测性和指标

信息中心

在 Google Cloud 控制台中的 Kubernetes 集群页面中，可观测性标签页显示 TPU 可观测性指标。如需了解详情，请参阅 GKE 可观测性指标。

仅当您在 GKE 集群中启用了系统指标时，系统才会填充 TPU 信息中心。

运行时指标

在 GKE 1.27.4-gke.900 版或更高版本中，使用 JAX 0.4.14 版或更高版本并指定 containerPort: 8431 的 TPU 工作负载会以 GKE 系统指标的形式，导出 TPU 利用率指标。Cloud Monitoring 提供以下指标来监控 TPU 工作负载的运行时性能：

工作周期：TensorCore 在 TPU 芯片上活跃处理的时间占过去的采样周期（60 秒）的百分比。百分比越大，表示 TPU 利用率越高。
已使用的内存：已分配的加速器内存量（以字节为单位）。每 60 秒采样一次。
内存总量：加速器内存总量（以字节为单位）。每 60 秒采样一次。

这些指标位于 Kubernetes 节点 (k8s_node) 和 Kubernetes 容器 (k8s_container) 架构中。

Kubernetes 容器：

kubernetes.io/container/accelerator/duty_cycle
kubernetes.io/container/accelerator/memory_used
kubernetes.io/container/accelerator/memory_total

Kubernetes 节点：

kubernetes.io/node/accelerator/duty_cycle
kubernetes.io/node/accelerator/memory_used
kubernetes.io/node/accelerator/memory_total

主机指标

在 GKE 1.28.1-gke.1066000 版或更高版本中，TPU 切片中的虚拟机会将 TPU 利用率指标导出为 GKE 系统指标。Cloud Monitoring 提供以下指标来监控 TPU 主机的性能：

TensorCore 利用率：当前使用的 TensorCore 百分比。TensorCore 值等于矩阵乘法单位 (MXU) 加上向量单元的总和。TensorCore 利用率值是过去采样周期（60 秒）内执行的 TensorCore 操作数量除以同一周期内支持的 TensorCore 操作数量。值越大，表示利用率越高。
内存带宽利用率：正在使用的加速器内存带宽的当前百分比。计算方法是将采样周期（60 秒）内使用的内存带宽除以同一采样周期内支持的最大带宽。

这些指标位于 Kubernetes 节点 (k8s_node) 和 Kubernetes 容器 (k8s_container) 架构中。

Kubernetes 容器：

kubernetes.io/container/accelerator/tensorcore_utilization
kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes 节点：

kubernetes.io/container/node/tensorcore_utilization
kubernetes.io/container/node/memory_bandwidth_utilization

如需了解详情，请参阅 Kubernetes 指标和 GKE 系统指标。

已知问题

在新的 TPU 分片节点报告可用 TPU 之前，集群自动扩缩器可能会错误地计算这些节点的容量。然后，集群自动扩缩器可能会执行额外的扩容，因此创建超出需要的节点。集群自动扩缩器会在常规缩容操作后缩容额外的节点（如果不需要）。
集群自动扩缩器会取消对等待时间超过 10 小时的 TPU 切片节点池的扩容操作。集群自动扩缩器稍后将重试此类扩容操作。对于不使用预留的客户，此行为可能会降低 TPU 可获取性。
容忍 TPU 污点的非 TPU 工作负载可能会阻止缩容在排空 TPU 切片节点池期间重新创建的节点池。
内存带宽利用率指标不适用于 v5e TPU。

在 GKE Standard 中部署 TPU 工作负载

须知事项

GKE 中的 TPU 可用性

规划 TPU 配置

确保为按需或 Spot 虚拟机提供足够的配额

确保预留可用性

创建集群

创建节点池

单主机 TPU 切片

gcloud

Terraform

控制台

多主机 TPU 切片

gcloud

Terraform

控制台

预配状态

在 TPU 切片节点上运行工作负载

工作负载准备

示例 1：运行显示 TPU 切片节点池中可用 TPU 芯片数量的工作负载

示例 2：运行显示 TPU 切片中可用 TPU 芯片数量的工作负载

使用加速器（GPU 和 TPU）升级节点池

清理

其他配置

多切片

迁移 TPU 预留

日志记录

使用 GKE 节点自动预配功能

TPU 切片节点自动修复

配置 TPU 分片节点的正常终止

在不使用特权模式的情况下运行容器

可观测性和指标

信息中心

运行时指标

主机指标

已知问题

后续步骤