将 TPU 与 KubeRay 结合使用#

本文档提供了将 TPU 与 KubeRay 结合使用的技巧。

TPUs 可在 Google Kubernetes Engine (GKE) 上使用。要在 Kubernetes 中使用 TPU，需要配置 Kubernetes 设置，并在 RayCluster CR 配置中添加额外的值。请参考启动 Google Cloud GKE 集群以使用 KubeRay 搭配 TPU 来配置 GKE 上的 TPU。

关于 TPU#

TPUs 是定制设计的 AI 加速器，针对大型 AI 模型的训练和推理进行了优化。TPU 主机是运行在连接到 TPU 硬件的物理计算机上的虚拟机。TPU 工作负载可以在一个或多个主机上运行。TPU Pod slice 是芯片的集合，这些芯片物理上共置并通过高速芯片间互连 (ICI) 连接。单主机 TPU Pod slice 包含独立的 TPU VM 主机，并通过数据中心网络 (DCN) 而不是 ICI 互连进行通信。多主机 TPU Pod slice 包含两个或多个互连的 TPU VM 主机。在 GKE 中，多主机 TPU Pod slice 运行在自己的节点池上，GKE 以节点池为单位（而不是以单个节点为单位）对其进行原子缩放。Ray 支持单主机和多主机 TPU Pod slice 平滑地扩展到多个 slice，从而实现更大的并行性以支持更大的工作负载。

快速入门：在 GKE 上使用 TPU 服务 Stable Diffusion 模型#

在设置好带有 TPU 和 Ray TPU 初始化 Webhook 的 GKE 集群后，您可以在 Ray 上运行带有 TPU 的工作负载。在 GKE 上使用 TPU 服务 Stable Diffusion 模型展示了如何在单主机 TPU 上使用 KubeRay 服务模型。

为 Ray Pod 配置 TPU 使用#

使用任何 TPU 加速器都需要在 RayCluster 的 workerGroupSpecs 的 container 字段中指定 google.com/tpu 资源 limits 和 requests。此资源指定了 GKE 为每个 Pod 分配的 TPU 芯片数量。KubeRay v1.1.0 为 RayCluster custom resource 添加了一个 numOfHosts 字段，用于指定每个 worker group replica 要创建的 TPU 主机数量。对于多主机 worker groups，Ray 将 replica 视为 Pod slice 而不是单个 worker，并为每个 replica 创建 numOfHosts worker 节点。此外，GKE 使用 gke-tpu 节点选择器将 TPU Pod 调度到与所需 TPU 加速器和拓扑匹配的节点上。

以下是一个 RayCluster worker group 的配置片段，其中包含 2 个 Ray TPU worker Pod。Ray 将每个 worker 调度到属于同一 TPU Pod slice 的各自 GKE v4 TPU 节点上。

   groupName: tpu-group
   replicas: 1
   minReplicas: 0
   maxReplicas: 1
   numOfHosts: 2
   ...
   template:
       spec:
        ...
        containers:
         - name: ray-worker
           image: rayproject/ray:2.9.0-py310
           ...
           resources:
             google.com/tpu: "4" # Required to use TPUs.
             ...
           limits:
             google.com/tpu: "4" # The resources and limits value is expected to be equal.
             ...
        nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
            cloud.google.com/gke-tpu-topology: 2x2x2
            ...

TPU 工作负载调度#

Ray 部署具有 TPU 资源的 Ray Pod 后，该 Ray Pod 可以执行使用 TPU 请求注解的任务和 Actor。Ray 支持将 TPU 作为自定义资源。任务或 Actor 使用装饰器 @ray.remote(resources={"TPU": NUM_TPUS}) 请求 TPU。

多主机 TPU 自动扩缩容#

KubeRay 1.1.0 及更高版本以及 Ray 2.32.0 及更高版本支持多主机 TPU 自动扩缩容。Ray multi-host TPU worker group 是指定了“google.com/tpu” Kubernetes 容器限制或请求并且 numOfHosts 大于 1 的 worker group。Ray 将 Ray multi-host TPU worker group 的每个 replica 视为一个 TPU Pod slice 并对其进行原子缩放。向上扩容时，multi-host worker groups 为每个 replica 创建 numOfHosts 个 Ray worker。同样，Ray 以 numOfHosts 个 worker 为单位缩小 multi-host worker group 的 replica。当 Ray 计划删除 multi-host TPU worker group 中的单个 Ray worker 时，它会终止该 worker 所属的整个 replica。在 multi-host worker groups 上调度 TPU 工作负载时，请确保 Ray 任务或 Actor 在 worker group replica 中的每个 TPU VM 主机上运行，以避免 Ray 缩减空闲的 TPU worker。