在 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.requestsresources.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 芯片数量。对于 limitsrequests,值必须相同。
  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,而无需排入队列来实现可用性。