本页面介绍了如何通过在 Google Kubernetes Engine (GKE) Autopilot 集群中使用 Cloud TPU 加速器 (TPU) 来加速机器学习 (ML) 工作负载。该指导可帮助您为机器学习应用框架选择正确的库,将 TPU 工作负载设置为在 GKE 上以最佳方式运行,并在部署后监控工作负载。
本页面适用于希望在 TPU 上准备和运行机器学习工作负载的平台管理员和运维人员、数据和 AI 专家以及应用开发者。如需详细了解我们在 Google Cloud 内容中提及的常见角色、职责和示例任务,请参阅常见的 GKE Enterprise 用户角色和任务。
在阅读本页面之前,请确保您熟悉以下资源:
Autopilot 中 TPU 的工作原理
如需在 Autopilot 工作负载中使用 TPU,请在工作负载清单中指定以下各项:
spec.nodeSelector
字段中的 TPU 版本。spec.nodeSelector
字段中的 TPU 拓扑。指定的 TPU 版本必须支持该拓扑。spec.containers.resources.requests
和spec.containers.resources.limits
字段中的 TPU 芯片数量。
部署工作负载时,GKE 会预配具有所请求 TPU 配置的节点,并在这些节点上调度 Pod。GKE 将每个工作负载放置在其自己的节点上,以便每个 Pod 都可以访问相应节点的全部资源,并将中断风险降至最低。
Autopilot 中的 TPU 与以下功能兼容:
规划 TPU 配置
在使用本指南部署 TPU 工作负载之前,请根据您的模型及其所需的内存量来规划 TPU 配置。如需了解详情,请参阅规划 TPU 配置。
价格
如需了解价格信息,请参阅 Autopilot 价格。
准备工作
在开始之前,请确保您已执行以下任务:
- 启用 Google Kubernetes Engine API。 启用 Google Kubernetes Engine API
- 如果您要使用 Google Cloud CLI 执行此任务,请安装并初始化 gcloud CLI。 如果您之前安装了 gcloud CLI,请运行
gcloud components update
以获取最新版本。
- 确保您拥有运行 GKE 1.29.2-gke.1521000 版或更高版本的 Autopilot 集群。
- 如需使用预留的 TPU,请确保您已有特定容量预留。如需查看相关说明,请参阅使用预留的可用区级资源。
确保 TPU 和其他 GKE 资源的配额
以下部分可帮助您确保使用 GKE 中的 TPU 时有足够的配额。 如需创建 TPU 切片节点,您必须拥有可用的 TPU 配额,除非您使用现有容量预留。如果您使用的是预留的 TPU,请跳过此部分。在 GKE 中创建 TPU 切片节点需要 Compute Engine API 配额 (compute.googleapis.com),而不是 Cloud TPU API 配额 (tpu.googleapis.com)。配额的名称在常规 Autopilot Pod 和 Spot Pod 中有所不同。
如需查看 TPU 的 Compute Engine API 配额的限制和当前用量,请按照以下步骤操作:
转到 Google Cloud 控制台中的配额页面。
在
过滤条件框中,执行以下操作:选择服务属性,输入 Compute Engine API,然后按 Enter 键。
选择类型属性,然后选择配额。
选择名称属性,然后根据
cloud.google.com/gke-tpu-accelerator
节点选择器中的 TPU 版本和值输入配额的名称。例如,如果您计划创建在cloud.google.com/gke-tpu-accelerator
节点选择器中其值为tpu-v5-lite-podslice
的按需 TPU v5e 节点,请输入TPU v5 Lite PodSlice chips
。TPU 版本 cloud.google.com/gke-tpu-accelerator
按需实例的配额名称 Spot2 实例的配额名称 TPU v3 tpu-v3-device
TPU v3 Device chips
Preemptible TPU v3 Device chips
TPU v3 tpu-v3-slice
TPU v3 PodSlice chips
Preemptible TPU v3 PodSlice chips
TPU v4 tpu-v4-podslice
TPU v4 PodSlice chips
Preemptible TPU v4 PodSlice chips
TPU v5e tpu-v5-lite-podslice
TPU v5 Lite PodSlice chips
Preemptible TPU v5 Lite PodSlice chips
TPU v5p tpu-v5p-slice
TPU v5p chips
Preemptible TPU v5p chips
TPU Trillium tpu-v6e-slice
TPU v6e Slice chips
Preemptible TPU v6e Lite PodSlice chips
选择维度(例如位置)属性,然后输入
region:
,后跟您计划在其中创建 GKE 中 TPU 的区域的名称。例如,如果您计划在可用区us-west4-a
中创建 TPU 切片节点,请输入region:us-west4
。TPU 配额是区域性的,因此同一区域内的所有可用区都使用同一 TPU 配额。
如果没有配额与您输入的过滤条件匹配,则说明项目尚未获得所需区域的任何指定配额,您必须申请增加 TPU 配额。
创建 TPU 预留后,相应配额的限制值和当前使用值都会增加 TPU 预留中的芯片数量。例如,为在 cloud.google.com/gke-tpu-accelerator
节点选择器中其值为 tpu-v5-lite-podslice
的 16 个 TPU v5e 芯片创建预留后,相关区域中 TPU v5 Lite PodSlice chips
配额的限制和当前用量都会增加 16。
其他 GKE 资源的配额
您可能需要在 GKE 创建资源的区域增加以下与 GKE 相关的配额。
- Persistent Disk SSD (GB) 配额:默认情况下,每个 Kubernetes 节点的启动磁盘需要 100GB。因此,此配额应至少设置为您预计创建的 GKE 节点数上限和 100GB 的乘积(节点数 * 100GB)。
- 使用中的 IP 地址配额:每个 Kubernetes 节点使用一个 IP 地址。因此,此配额应至少设置为您预计创建的 GKE 节点数上限。
- 确保
max-pods-per-node
与子网范围一致:每个 Kubernetes 节点都对 Pod 使用次要 IP 范围。例如,max-pods-per-node
为 32 时需要 64 个 IP 地址,这相当于每个节点一个 /26 子网。请注意,此范围不应与任何其他集群共享。为避免耗尽 IP 地址范围,请使用--max-pods-per-node
标志来限制允许在节点上调度的 Pod 数量。max-pods-per-node
的配额应至少设置为您预计创建的 GKE 节点数上限。
如需申请增加配额,请参阅申请更高配额。
在 GKE 中预配 TPU 的选项
借助 GKE Autopilot,您可以通过使用 Kubernetes nodeSelector 直接在各个工作负载中使用 TPU。
或者,您也可以使用自定义计算类来请求 TPU。借助自定义计算类,平台管理员可以定义节点配置的层次结构,以便 GKE 在节点扩缩决策期间确定优先级,从而使工作负载在您选择的硬件上运行。
如需查看相关说明,请参阅使用自定义计算类集中预配 TPU 部分。
准备 TPU 应用
TPU 工作负载具有以下准备要求。
- JAX、PyTorch 和 TensorFlow 等框架使用
libtpu
共享库访问 TPU 虚拟机。libtpu
包括 XLA 编译器、TPU 运行时软件和 TPU 驱动程序。每个 PyTorch 和 JAX 版本都需要特定的libtpu.so
版本。如需使用 GKE 中的 TPU,请确保您使用以下版本:TPU 类型 libtpu.so
版本TPU Trillium (v6e)
tpu-v6e-slice
- 推荐的 jax[tpu] 版本:v0.4.9 或更高版本
- 推荐的 torchxla[tpuvm] 版本:v2.1.0 或更高版本。
TPU v5e
tpu-v5-lite-podslice
- 推荐的 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 v3
tpu-v3-slice
tpu-v3-device
- 推荐的 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
排序且为零索引。
在
completionMode: Indexed
,subdomain
,parallelism > 1
情况下创建 Job 并请求google.com/tpu
属性时,GKE 会使用变更 webhook 自动注入这些环境变量。GKE 会添加一个无头 Service,以便为支持该 Service 的 Pod 添加 DNS 记录。
完成工作负载准备后,您可以运行使用 TPU 的 Job。
在工作负载中请求 TPU
本部分介绍了如何创建在 Autopilot 中请求 TPU 的 Job。在任何需要 TPU 的工作负载中,您必须指定以下各项:
- 用于 TPU 版本和拓扑的节点选择器
- 工作负载中容器的 TPU 芯片数量
如需查看支持的 TPU 版本、拓扑以及切片中对应的 TPU 芯片数量和节点数量列表,请参阅选择 Autopilot TPU 配置。
有关工作负载中的 TPU 请求的注意事项
Pod 中只有一个容器可以使用 TPU。容器请求的 TPU 芯片数量必须等于连接到切片中节点的 TPU 芯片数量。例如,如果您请求具有 2x4
拓扑的 TPU v5e (tpu-v5-lite-podslice
),则可以请求以下任一项:
4
个芯片,会创建两个分别具有 4 个 TPU 芯片的多主机节点8
个芯片,会创建一个具有 8 个 TPU 芯片的单主机节点
为最大限度地提高成本效益,最佳实践是始终使用您请求的切片中的所有 TPU 芯片。如果您请求一个包含两个节点(每个节点具有 4 个 TPU 芯片)的多主机切片,则应该部署一个在两个节点上运行并使用切片中全部 8 个 TPU 芯片的工作负载。
创建请求 TPU 的工作负载
以下步骤会创建一个请求 TPU 的 Job。如果您的工作负载在多主机 TPU 切片上运行,则还必须创建一个按名称选择工作负载的无头 Service。该无头 Service 通过更新 Kubernetes DNS 配置以指向工作负载中的 Pod,可让多主机切片中不同节点上的 Pod 相互通信。
将以下清单保存为
tpu-autopilot.yaml
:apiVersion: v1 kind: Service metadata: name: headless-svc spec: clusterIP: None selector: job-name: tpu-job --- apiVersion: batch/v1 kind: Job metadata: name: tpu-job spec: backoffLimit: 0 completions: 4 parallelism: 4 completionMode: Indexed template: spec: subdomain: headless-svc restartPolicy: Never nodeSelector: cloud.google.com/gke-tpu-accelerator: TPU_TYPE cloud.google.com/gke-tpu-topology: TOPOLOGY 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. 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: NUMBER_OF_CHIPS limits: cpu: 10 memory: 500Gi google.com/tpu: NUMBER_OF_CHIPS
替换以下内容:
TPU_TYPE
:要使用的 TPU 类型,例如tpu-v4-podslice
。必须是 GKE 支持的值。TOPOLOGY
:切片中 TPU 芯片的排列方式,例如2x2x4
。必须是所选 TPU 类型支持的拓扑。NUMBER_OF_CHIPS
:容器要使用的 TPU 芯片数量。对于limits
和requests
,值必须相同。
部署作业:
kubectl create -f tpu-autopilot.yaml
创建此作业时,GKE 会自动执行以下操作:
- 预配节点以运行 Pod。这些节点可以是单主机切片或多主机切片,具体取决于您指定的 TPU 类型、拓扑和资源请求。
- 向 Pod 添加污点并向节点添加容忍设置,以防止任何其他工作负载与 TPU 工作负载在相同节点上运行。
创建请求 TPU 和集合调度的作业
在 TPU Trillium 中,您可以使用集合调度对 TPU 切片节点进行分组。对这些 TPU 切片节点进行分组后,您可以更轻松地调整副本数量,以满足工作负载需求。 Google Cloud 会控制软件更新,以确保集合中的切片充足,始终可用于处理流量。
TPU Trillium 支持为运行推理工作负载的单主机和多主机节点池提供集合调度。以下内容介绍了集合调度行为如何取决于您使用的 TPU 切片类型:
- 多主机 TPU 切片:GKE 会将多主机 TPU 切片分组以形成集合。每个 GKE 节点池都是此集合中的副本。如需定义集合,请创建多主机 TPU 切片并为集合分配唯一名称。如需向集合中添加更多 TPU 切片,请创建另一个具有相同集合名称和工作负载类型的多主机 TPU 切片节点池。
- 单主机 TPU 切片:GKE 会将整个单主机 TPU 切片节点池视为一个集合。如需向集合中添加更多 TPU 切片,您可以调整单主机 TPU 切片节点池的大小。
如需了解集合调度的限制,请参阅集合调度的工作原理
使用多主机 TPU 切片
多主机 TPU 切片节点中的集合调度适用于 1.31.2-gke.1537000 版及更高版本的 Autopilot 集群。只有 1.31.2-gke.1115000 或更高版本支持具有 2x4
拓扑的多主机 TPU 切片节点。如需创建多主机 TPU 切片节点并将其作为集合进行分组,请在工作负载规范中添加以下 Kubernetes 标签:
cloud.google.com/gke-nodepool-group-name
:每个集合在集群级别应具有唯一的名称。cloud.google.com/gke-nodepool-group-name
标签中的值必须符合集群标签的要求。cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
例如,以下代码块定义了一个包含多主机 TPU 切片的集合:
nodeSelector: cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME} cloud.google.com/gke-workload-type: HIGH_AVAILABILITY cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice cloud.google.com/gke-tpu-topology: 4x4 ...
使用单主机 TPU 切片
单主机 TPU 切片节点中的集合调度适用于 1.31.2-gke.1088000 版及更高版本的 Autopilot 集群。如需创建单主机 TPU 切片节点并将其作为集合进行分组,请在工作负载规范中添加 cloud.google.com/gke-workload-type:HIGH_AVAILABILITY
标签。
例如,以下代码块定义了一个包含单主机 TPU 切片的集合:
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 2x2
cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
...
使用自定义计算类集中预配 TPU
如需使用自定义计算类预配 TPU,请执行以下操作:
确保您的集群具有可用的自定义计算类,以便选择 TPU。如需了解如何在自定义计算类中指定 TPU,请参阅 TPU 规则。
将以下清单保存为
tpu-job.yaml
:apiVersion: v1 kind: Service metadata: name: headless-svc spec: clusterIP: None selector: job-name: tpu-job --- apiVersion: batch/v1 kind: Job metadata: name: tpu-job spec: backoffLimit: 0 completions: 4 parallelism: 4 completionMode: Indexed template: spec: subdomain: headless-svc restartPolicy: Never nodeSelector: cloud.google.com/compute-class: TPU_CLASS_NAME 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. 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: NUMBER_OF_CHIPS limits: cpu: 10 memory: 500Gi google.com/tpu: NUMBER_OF_CHIPS
替换以下内容:
TPU_CLASS_NAME
:指定 TPU 的现有自定义计算类的名称。NUMBER_OF_CHIPS
:容器要使用的 TPU 芯片数量。对于limits
和requests
,值必须相同,并且与所选自定义计算类的tpu.count
字段中的值相同。
部署作业:
kubectl create -f tpu-workload.yaml
创建此作业时,GKE 会自动执行以下操作:
- 预配节点以运行 Pod。这些节点可以是单主机切片或多主机切片,具体取决于您指定的 TPU 类型、拓扑和资源请求。 根据最高优先级的 TPU 资源的可用性,GKE 可能会回退到较低优先级,以最大限度地提高可获取性。
- 向 Pod 添加污点并向节点添加容忍设置,以防止任何其他工作负载与 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. 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
示例:显示单个节点中的 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. 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
观察和监控 TPU
信息中心
在 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 系统指标。
日志记录
GKE 节点(包括 TPU 虚拟机)上运行的容器发出的日志会由 GKE 日志记录代理收集,发送到 Logging,并且显示在 Logging 中。
针对 Autopilot 中的 TPU 工作负载的建议
以下建议可能会提高 TPU 工作负载的效率:
- 在长达七天宽限期内使用延长运行时间的 Pod,然后 GKE 终止 Pod 进行缩容或节点升级。您可以将维护窗口和排除项与延长运行时间的 Pod 搭配使用,以进一步延迟自动节点升级。
- 使用容量预留可确保工作负载接收请求的 TPU,而无需排入队列来实现可用性。