在 GKE Autopilot 上部署 TPU 工作负载

本页面介绍了如何通过在 Google Kubernetes Engine (GKE) Autopilot 集群中使用 Cloud TPU 加速器 (TPU) 进行部署，加速机器学习 (ML) 工作负载。您应该已经熟悉以下概念：

  1. Cloud TPU 简介
  2. Cloud TPU 系统架构
  3. 关于 GKE 中的 TPU

Autopilot 中 TPU 的工作原理

如需在 Autopilot 工作负载中使用 TPU，请在工作负载清单中请求 TPU 版本和该 TPU 版本支持的拓扑。然后，您可以使用 Kubernetes resources.requests 和 resources.limits 字段指定工作负载要使用的 TPU 芯片的数量。部署工作负载时，GKE 会预配具有所请求 TPU 配置的节点，并在这些节点上调度 Pod。GKE 将每个工作负载放置在其自己的节点上，以便每个 Pod 都可以访问相应节点的全部资源，并将中断风险降至最低。

Autopilot 中的 TPU 与以下功能兼容：

  1. Spot Pod
  2. 特定容量预留
  3. 延长运行时间的 Pod

规划 TPU 配置

在请求 TPU 之前，请根据工作负载 CPU 和内存要求确定所需的配置。您必须确定以下各项：

  • TPU 版本：特定 Cloud TPU 版本，例如 v5e。
  • 所选 TPU 版本的拓扑：TPU 的排列方式和数量。

您选择的 TPU 版本和拓扑决定了 GKE 将节点预配为单主机切片还是多主机切片。在单主机切片中，每个节点都独立于其他 TPU 切片节点。在多主机切片中，GKE 会创建一组节点，这些节点在 TPU 切片中具有互连虚拟机。多主机切片具有原子性，这意味着 GKE 将整个互连节点组作为一个单元进行扩缩。

如需了解可用的 TPU 版本、对应的拓扑、CPU 和内存容量以及生成的切片类型，请参阅选择 Autopilot 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 配额

如需创建 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 配额的限制和当前用量，请按照以下步骤操作：

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

    转到“配额”

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

    1. 选择服务属性，输入 Compute Engine API，然后按 Enter 键。

    2. 选择类型属性，然后选择配额。

    3. 选择名称属性，然后根据所需的 TPU 类型输入配额名称，如下所示：

      • TPU v5p (tpu-v5p-slice)：TPU v5p 芯片
      • TPU v5e (tpu-v5-lite-podslice)：TPU v5 Lite PodSlice 芯片
      • TPU v5e (tpu-v5-lite-device)：TPU v5 Lite Device 芯片
      • TPU v4 (tpu-v4-podslice)：TPU v4 PodSlice 芯片

      对于 Spot Pod，请选择相应的“可抢占”配额。

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

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

准备 TPU 应用

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

  1. 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
    TPU v5p
    tpu-v5p-slice
    • 推荐的 jax[tpu] 版本：0.4.19 或更高版本。
    • 推荐的 torchxla[tpuvm] 版本：建议使用 2023 年 10 月 23 日的夜间版本 build。
    TPU v4
    tpu-v4-podslice
  2. 为请求 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 相互通信。

  1. 将以下清单保存为 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，值必须相同。

  2. 部署作业：

    kubectl create -f tpu-autopilot.yaml

创建此作业时，GKE 会自动执行以下操作：

  1. 预配节点以运行 Pod。这些节点可以是单主机切片或多主机切片，具体取决于您指定的 TPU 类型、拓扑和资源请求。
  2. 向 Pod 添加污点并向节点添加容忍设置，以防止任何其他工作负载与 TPU 工作负载在相同节点上运行。

示例：显示多主机切片中的 TPU 芯片总数

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

  • TPU 版本：TPU v4
  • 拓扑：2x2x4

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

  1. 将以下清单保存为 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
  2. 部署清单： 
    kubectl create -f available-chips-multihost.yaml

    GKE 会运行具有四个虚拟机的 TPU v4 切片（多主机 TPU 切片）。该切片具有 16 个互连 TPU 芯片。

  3. 验证作业是否创建了四个 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
  4. 获取其中一个 Pod 的日志： 
    kubectl logs POD_NAME

    将 POD_NAME 替换为创建的某个 Pod 的名称。例如 tpu-job-podslice-0-5cd8r。

    输出类似于以下内容： 

    TPU cores: 16

示例：显示单个节点中的 TPU 芯片数量

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

  • TPU 版本：TPU v5e
  • 拓扑：2x4

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

  1. 将以下清单保存为 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
  2. 部署清单： 
    kubectl create -f available-chips-singlehost.yaml

    GKE 会预配具有八个单主机 TPU 切片的节点，这些切片使用 TPU v5e。每个 TPU 节点都具有八个 TPU 芯片（单主机 TPU 切片）。

  3. 获取相应 Pod 的日志： 
    kubectl logs tpu-job-jax-v5

    输出类似于以下内容： 

    Total TPU chips: 8

可观测性和指标

信息中心

在 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，而无需排入队列来实现可用性。