在 KubeRay 中使用 TPU#

本文档提供了使用 KubeRay 进行 TPU 使用的技巧。

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

关于 TPU#

TPU 是专门设计的 AI 加速器，针对大型 AI 模型的训练和推理进行了优化。TPU 主机是在连接到 TPU 硬件的物理计算机上运行的虚拟机。TPU 工作负载可以在一个或多个主机上运行。TPU Pod 切片是所有物理上共存并通过高速芯片间互连（ICI）连接的芯片集合。单主机 TPU Pod 切片包含独立的 TPU 虚拟机主机，并通过数据中心网络（DCN）而不是 ICI 互连进行通信。多主机 TPU Pod 切片包含两个或更多互连的 TPU 虚拟机主机。在 GKE 中，多主机 TPU Pod 切片运行在自己的节点池上，GKE 通过节点池而不是单个节点来原子化地扩展它们。Ray 使得单主机和多主机 TPU Pod 切片能够无缝扩展到多个切片，从而支持更大的工作负载，实现更高的并行性。

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

在设置了带有TPU的GKE集群和Ray TPU初始化webhook之后，在Ray上运行带有TPU的工作负载。在GKE上使用TPU提供Stable Diffusion模型展示了如何在单主机TPU上使用KubeRay提供模型。

为 TPU 使用配置 Ray Pod#

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

以下是一个配置片段，用于包含2个Ray TPU工作Pod的RayCluster工作组。Ray将每个工作节点调度到属于同一TPU Pod切片的自己的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 请求的任务和角色。Ray 支持 TPU 作为 自定义资源。任务或角色使用装饰器 @ray.remote(resources={"TPU": NUM_TPUS}) 请求 TPU。

多主机 TPU 自动扩展#

在 Kuberay 1.1.0 及以上版本和 Ray 2.32.0 及以上版本中支持多主机 TPU 自动扩展。Ray 多主机 TPU 工作组是指定了 “google.com/tpu” Kubernetes 容器限制或请求并且 NumOfHosts 大于 1 的工作组。Ray 将 Ray 多主机 TPU 工作组的每个副本视为一个 TPU Pod 切片，并以原子方式扩展它们。在扩展时，多主机工作组为每个副本创建 NumOfHosts 个 Ray 工作器。同样，Ray 通过 NumOfHosts 个工作器缩减多主机工作组的副本。当 Ray 调度删除多主机 TPU 工作组中的单个 Ray 工作器时，它会终止该工作器所属的整个副本。在多主机工作组上调度 TPU 工作负载时，确保 Ray 任务或角色在每个工作组副本中的每个 TPU VM 主机上运行，以避免 Ray 缩减空闲的 TPU 工作器。

进一步的参考和讨论#