在 KubeRay 中使用 TPU#

本文档提供了在 KubeRay 中使用 TPU 的技巧。

TPU 在 Google Kubernetes Engine (GKE) 上可用。要将 TPU 与 Kubernetes 结合使用，请配置 Kubernetes 设置并为 RayCluster CR 配置添加额外的值。通过参考 为 KubeRay 设置带 TPU 的 Google Cloud GKE 集群 来配置 GKE 上的 TPU。

关于 TPU#

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

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

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

为 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 请求注释的任务和 Actor。Ray 支持将 TPU 作为自定义资源。任务或 Actor 使用装饰器 @ray.remote(resources={"TPU": NUM_TPUS}) 请求 TPU。

TPU 默认标签#

在 Google Cloud TPU 上使用 KubeRay 时，Ray 会自动检测并添加以下标签来描述底层计算。这些对于调度必须跨越整个 TPU “切片”（一组互连的主机）的分布式工作负载至关重要。

  • ray.io/accelerator-type：TPU 加速器的类型，例如 TPU-V6E。

  • ray.io/tpu-slice-name：TPU Pod 或切片的名称。Ray 使用此标签来确保作业的所有工作节点都位于*同一*切片上。

  • ray.io/tpu-worker-id：切片内的整数工作节点 ID。

  • ray.io/tpu-topology：切片的物理拓扑。

  • ray.io/tpu-pod-type：TPU Pod 类型，定义了大小和 TPU 代，例如 v4-8 或 v5p-16。

您可以使用这些标签来调度一个请求整个 TPU 切片的 placement_group。例如，请求 v6e-16 切片上的所有 TPU 设备

# Request 4 bundles, one for each TPU VM in the v6e-16 slice.
pg = placement_group(
    [{"TPU": 4}] * 4,
    strategy="SPREAD",
    bundle_label_selector=[{
        "ray.io/tpu-pod-type": "v6e-16"
    }] * 4
)
ray.get(pg.ready())

TPU 调度实用库#

ray.util.tpu 包引入了许多与调度相关的 TPU 实用程序，它们简化了在多主机和/或多切片配置中使用 TPU 的过程。这些实用程序利用在 Google Kubernetes Engine (GKE) 上使用 KubeRay 运行时设置的默认 Ray 节点标签。

以前，当通过多个任务或 Actor 简单地请求 TPU，例如使用 resources={"TPU": 4} 时，从不能保证所调度的 Ray 节点属于同一切片，甚至属于同一 TPU 代。为了解决后者问题，Ray 引入了基于标签的调度 API和默认标签（如 ray.io/accelerator-type）来描述底层计算。

更进一步，新的 TPU 实用库利用默认节点标签和标签选择器 API 来抽象 TPU 调度的复杂性，尤其是在多主机切片方面。核心抽象是 SlicePlacementGroup。

SlicePlacementGroup#

SlicePlacementGroup 类提供了一个高级接口，用于预留一个或多个完整、可用的 TPU 切片，并创建一个 Ray Placement Group，该组限制在这些切片内。这确保了放置组内的所有 bundle（以及在它们上调度的任务/Actor）都在属于已预留物理 TPU 切片的工作节点上运行。

工作原理

  1. 预留： 当您创建 SlicePlacementGroup 时，它首先与 Ray 调度器交互，以查找符合您指定的 topology 和 accelerator_version 的可用 TPU 切片。它通过使用一个小型内部放置组来临时预留切片的“head”TPU 节点（工作节点 ID 0）来实现这一点。

  2. 切片标识： 从预留的 head 节点，它检索唯一的切片名称（使用 ray.io/tpu-slice-name 默认标签）。此标签是通过 GKE Webhook 注入的环境变量设置的。GKE Webhook 还确保具有 numOfHosts > 1 的 KubeRay Pod 具有亲和性，并被调度到同一个 GKE 节点池中，该节点池与 TPU 多主机切片是一对一的关系。

  3. 主放置组创建： 然后，它创建您请求的主放置组。此组包含代表切片中每个主机（VM）的 bundle。对于每个切片，它使用 bundle_label_selector 来定位上一步中标识的特定 ray.io/tpu-slice-name。这确保了给定切片的所有 bundle 都位于该确切切片内的节点上。

  4. 句柄： 它返回一个 SlicePlacementGroup 句柄，该句柄公开了底层的 Ray PlacementGroup 对象（.placement_group），以及诸如工作节点数量（.num_workers）和每个主机芯片数量（.chips_per_host）等有用属性。

用法

您通常使用 slice_placement_group 函数创建 SlicePlacementGroup。

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.tpu import slice_placement_group

# Reserve two v6e TPU slices, each with a 4x4 topology (16 chips each).
# This topology typically has 4 VM workers, each with 4 chips.
slice_handle = slice_placement_group(topology="4x4", accelerator_version="v6e", num_slices=2)
slice_pg = slice_handle.placement_group

print("Waiting for placement group to be ready...")
ray.get(slice_pg.ready(), timeout=600) # Increased timeout for potential scaling
print("Placement group ready.")

@ray.remote(num_cpus=0, resources={"TPU": 4})
def spmd_task(world_size, rank):
    pod_name = ray.util.tpu.get_current_pod_name()
    chips_on_node = ray.util.tpu.get_num_tpu_chips_on_node()
    print(f"Worker Rank {rank}/{world_size}: Running on slice '{pod_name}' with {chips_on_node} chips.")
    return rank

# Launch one task per VM in the reserved slices. The num_workers field describes the total
# number of VMs across all slices in the SlicePlacementGroup.
tasks = [
    spmd_task.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=slice_pg,
        )
    ).remote(world_size=slice_handle.num_workers, rank=i)
    for i in range(slice_handle.num_workers)
]

results = ray.get(tasks)
print(f"Task results: {results}")

TPU Pod 信息实用程序#

这些函数提供有关给定工作节点所属的 TPU Pod 的信息。如果工作节点未运行在 TPU 上，它们将返回 None。

  • ray.util.tpu.get_current_pod_name() -> Optional[str]

返回工作节点所属的 TPU Pod 的名称。

  • ray.util.tpu.get_current_pod_worker_count() -> Optional[int]

计算工作节点所属的 TPU Pod 的关联工作节点数量。

  • ray.util.tpu.get_num_tpu_chips_on_node() -> int

返回当前节点上的 TPU 芯片总数。如果未找到，则返回 0。

多主机 TPU 自动扩展#

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

进一步参考和讨论#