在 GKE 上使用 Kueue 和 LeaderWorkerSet 实现 TPU 动态切分

1. 简介

在此 Codelab 中，您将学习如何使用 GKE 动态切分来优化 Cloud TPU 资源的利用率。动态切片是一项强大的功能，可让您将原始 TPU 预配与工作负载调度分离。

具体而言，您将了解以下两种关键模式：

  • 子切片：将预配的大型 TPU 块拆分为较小的隔离切片，以用于较小的工作负载。
  • 超级切片：将多个已配置的 TPU 块拼接在一起，形成一个更大的虚拟切片，用于大规模工作负载。

您将应用这些模式，使用 Kueue、LeaderWorkerSet (LWS) 和 Gateway API 为大语言模型 (Qwen 397B) 部署高性能的分离式服务 (Prefill/Decode 分离) 架构。

架构

下图展示了 TPU 动态切分和分离式服务设置的概要架构：

TPU 动态切分架构

您将执行的操作

  • 预配已启用 GKE Slice 控制器的 GKE 集群。
  • 创建配置为增量配置的 GKE TPU 节点池。
  • 部署 Kueue 和 LeaderWorkerSet 以管理 TPU 工作负载。
  • 运行子切片工作负载，以验证 JAX TPU 在较小切片上的访问权限。
  • 运行超级切片工作负载，以验证 JAX TPU 在多个组合节点池中的访问权限。
  • 部署分离式服务设置，其中预填充和解码阶段在由 LLM 路由器协调的单独动态分配的 TPU 切片上运行。

所需条件

  • 网络浏览器，例如 Chrome。
  • 启用了结算功能的 Google Cloud 项目。
  • 重要提示：对 Cloud TPU7x (Ironwood) 全容量模式预留的访问权限。

2. 准备工作

创建或选择 Google Cloud 项目

创建 Google Cloud 项目

  1. 在 Google Cloud 控制台的项目选择器页面上，选择或创建一个 Google Cloud 项目。
  2. 确保您的 Cloud 项目已启用结算功能。了解如何检查项目是否已启用结算功能。

启动 Cloud Shell

Cloud Shell 是在 Google Cloud 中运行的命令行环境，预加载了必要的工具。

  1. 点击 Google Cloud 控制台顶部的激活 Cloud Shell。
  2. 连接到 Cloud Shell 后，验证您的身份验证：
    gcloud auth list
  3. 确认您的项目已配置：
    gcloud config get project
  4. 如果项目未按预期设置，请进行设置：
    export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

克隆演示版代码库

克隆包含此 Codelab 的清单和辅助脚本的代码库：

git clone --depth 1 --sparse https://github.com/GoogleCloudPlatform/devrel-demos.git
cd devrel-demos
git sparse-checkout set ai-ml/dynamic-slicing
cd ai-ml/dynamic-slicing

3. 配置环境

在预配资源之前，您需要配置环境变量。系统提供了一个帮助程序脚本 01_setup_env.sh 来生成 env.sh 文件。

运行设置脚本：

./01_setup_env.sh

系统会提示您输入多个值。按 [ENTER] 接受默认值，但请务必提供活动讲师提供的正确预订名称和预订块：

  • GCP 项目 ID：您当前的项目 ID。
  • GCP 项目编号：您的项目编号。
  • GKE 集群名称：tpu-serving-cluster（默认）。
  • TPU 节点池可用区：us-central1-ai1a（默认）。
  • Kubernetes 命名空间：llm-d-pd-disaggregation（默认）。
  • Cloud TPU 预留名称：[输入提供的预留名称]
  • Cloud TPU 预留块名称：block-0（默认）。
  • 权重对应的 GCS 存储分区名称：model-weights（默认值）。
  • TPU 机器类型：tpu7x-standard-4t（默认）。
  • Hugging Face 令牌：[如果需要，请输入您的 HF 令牌；如果使用预加载的权重，请按 Enter 键]

运行脚本后，将变量应用于当前会话：

source env.sh

4. 启用 API 和 AI Zone 功能

现在，您的环境已配置完毕，接下来需要启用必需的 Google Cloud API 和 AI Zone 可见性功能。我们提供了一个辅助脚本 02_enable_apis_and_features.sh。

运行脚本：

./02_enable_apis_and_features.sh

此脚本：

  1. 启用 GKE、Compute、IAM、Resource Manager、Filestore 和 Network Services API。
  2. 为 GKE 动态切片启用 ai-zones-visibility 预览版功能。

5. 预配 GKE 集群和 TPU 节点池

在此步骤中，您将预配底层网络基础架构、GKE 集群和 TPU 节点池。

TPU 节点池将配置为增量配置（使用 --placement-policy=superslice-policy 和 --reservation-affinity=specific），这会将每个节点池映射到 16 节点的原始 TPU 容量“立方体”（子块）。

运行配置脚本：

./03_create_cluster_and_nodes.sh

此脚本的作用：

  1. 创建 VPC 网络和子网：设置一个具有较大 MTU (8896) 的主 VPC 网络（针对 TPU 流量进行了优化）、一个 TPU 子网和一个 GKE Gateway 所需的代理专用子网。
  2. 创建 GKE 集群：预配启用了切片控制器 (--enable-slice-controller) 的标准 GKE 集群。
  3. 创建工作负载政策：定义一个名为 superslice-policy、类型为 HIGH_THROUGHPUT 且拓扑为 4x4x4 的资源政策。
  4. 创建 GKE TPU 节点池：预配两个节点池（tpu7-pool-1 和 tpu7-pool-2），每个节点池包含 16 个 tpu7x-standard-4t 节点。这表示两个独立的 16 节点立方体。

验证节点

脚本完成后，验证所有 32 个 TPU 节点是否已预配并注册：

kubectl get nodes -l google.com/tpu=present

您应该会在列表中看到 32 个节点。

6. 安装编排工具

动态切片依赖于多个 Kubernetes 控制器来协调作业和切片分配。您将安装：

  • JobSet：用于管理作业组（超切片必需）。
  • Kueue：用于排队、资源管理和拓扑感知调度 (TAS)。
  • LeaderWorkerSet (LWS)：用于管理复制的多节点 TPU 部署（LLM 服务所需）。
  • GKE Slice Controller（用户空间）：将 Kueue 与 TPU Cluster Director 连接起来，以动态管理物理切片。

运行安装脚本：

./04_install_kueue_lws_slice_controller.sh

验证 Slice Controller 是否已成功运行：

kubectl rollout status deployment/slice-controller-controller-manager -n slice-controller-system

7. 配置 Kueue 资源

现在，您需要定义表示 TPU 硬件拓扑的 Kueue 资源，并配置准入检查。

运行部署脚本：

./05_deploy_kueue_resources.sh

已部署的关键资源：

  • 拓扑 (slice-topology)：定义 Kueue 在调度时应考虑的 TPU 分区层次结构级别（从块到主机名）。
  • ResourceFlavor (slice-rf)：将 slice-topology 与 tpu7x 加速器相关联。
  • AdmissionCheck (ac)：配置 Kueue 以使用 GKE Slice Controller (accelerator.gke.io/slice) 在作业被准入时动态预配切片。
  • ClusterQueue (cq) 和 LocalQueue (lq)：设置工作负载将提交到的队列。
  • WorkloadPriorityClass（low-priority-1000、medium-priority-2000、high-priority-3000）：定义优先级，以实现抢占和基于优先级的调度。

验证资源：

kubectl get topology slice-topology
kubectl get resourceflavor slice-rf
kubectl get admissioncheck ac
kubectl get clusterqueue cq
kubectl get localqueue lq -n ${NAMESPACE}

8. 部署和验证使用细分功能的 TPU 访问权限

借助子切片，您可以在单个已配置的 TPU 块内运行多个较小的工作负载。在此步骤中，您将向由 4x4x4（64 个芯片 / 16 个虚拟机）块组成的集群提交请求 2x2x2 拓扑（8 个芯片 / 2 个虚拟机）的工作负载。

部署子切片工作负载：

./06_deploy_simple_subslicing.sh

此脚本应用 kueue-jobset-simple-subslicing.yaml。

工作原理：

  • JobSet 规范包含注释 cloud.google.com/gke-tpu-slice-topology: 2x2x2。
  • 它配置了 replicas: 6 和 parallelism: 2（完成次数：2）。这意味着 Kueue 将调度 6 个独立的作业，每个作业包含 2 个 pod。
  • 每个 pod 请求 google.com/tpu: "4"（1 个 TPU 虚拟机）。
  • Kueue 和 GKE Slice Controller 会动态划分 32 节点集群，以分配六个 2x2x2 slice。

验证 JAX 执行

监控 pod，直到它们开始运行：

kubectl get pods -n ${NAMESPACE} -l jobset.sigs.k8s.io/jobset-name=kueue-jobset-simple-subslicing

检查其中一个 pod 的日志，验证 JAX 是否已成功检测到其子切片上的 8 个 TPU 设备（核心）：

kubectl logs $(kubectl get pods -n ${NAMESPACE} -l jobset.sigs.k8s.io/jobset-name=kueue-jobset-simple-subslicing -o name | head -n 1) -n ${NAMESPACE}

您应该会看到指示以下内容的输出：Total TPU devices (cores): 8

9. 通过超切片部署和验证 TPU 访问权限

超级切片是一项强大的 GKE 功能，可让单个工作负载跨越多个物理 TPU 块（通常称为立方体或拓扑，例如 4x4x4）。通过将这些块拼接在一起，您可以形成更大的虚拟切片，用于大规模训练或服务工作负载。在此步骤中，您将部署一个请求 4x4x8 拓扑（128 个芯片 / 32 个虚拟机）的 JobSet。由于单个 4x4x4 块仅包含 64 个芯片，因此此工作负载超出了单个块的大小，需要 GKE 动态拼接 tpu7-pool-1 和 tpu7-pool-2 节点池以满足请求。

部署超级切片工作负载：

./07_deploy_simple_superslicing.sh

此脚本应用 kueue-jobset-simple-superslicing.yaml。

工作原理：

  • JobSet 模板包含注释 cloud.google.com/gke-tpu-slice-topology: 4x4x8。
  • 它会配置 parallelism: 32 和 completions: 32。
  • 每个 pod 请求 google.com/tpu: "4"。
  • 由于 4x4x8 拓扑需要所有 32 个节点，因此 Slice Controller 会动态配置 OCS（光路交换）网络，以将两个 16 节点池互连成一个 32 节点 ICI 网格。

验证 JobSet pod 是否成功运行，以及 JAX 是否检测到所有 128 个设备：

kubectl get pods -n ${NAMESPACE} -l jobset.sigs.k8s.io/jobset-name=kueue-jobset-simple-superslicing

查看其中一个 pod 的日志：

kubectl logs $(kubectl get pods -n ${NAMESPACE} -l jobset.sigs.k8s.io/jobset-name=kueue-jobset-simple-superslicing -o name | head -n 1) -n ${NAMESPACE}

您应该会看到 JAX 输出，其中显示了全局设备数量：Global device count: 128

10. 部署分离式服务（预填充/解码）

现在，您将使用 Prefill/Decode Disaggregation 部署端到端 LLM 服务堆栈。

在标准部署中，预填充（处理提示）和解码（生成令牌）在同一 TPU 上运行。由于预填充受计算限制，而解码受内存带宽限制，因此它们会发生冲突。分离式部署会在单独的 TPU 切片上运行它们，并通过网络传输 KV 缓存。

设置 LLM-D 和网关

设置命名空间、Hugging Face Secret 和 GKE 网关：

./08_setup_llm_d.sh

部署 LLM-D 路由器

部署将接收客户端请求并协调 Prefill 和 Decode 切片之间路由的路由器：

./09_deploy_llm_d_router.sh

部署预填充和解码工作负载

在动态分配的 TPU 切片上部署 vLLM 模型服务器：

./10_deploy_subslicing_pd_workload.sh

作用：

  • 部署 kueue-vllm-prefill-model-streamer（LWS 请求 2x2x2 TPU 切片）。
  • 部署 kueue-vllm-decode-model-streamer（LWS 请求 2x2x2 TPU 切片）。
  • 预填充切片会加载 Qwen 397B 模型权重，并充当 kv_producer。
  • 解码切片充当 kv_consumer。
  • 它们使用 TPUConnectorHMA 来传输 KV 缓存。

等待预填充 Pod 和解码 Pod 都开始运行：

kubectl get pods -n ${NAMESPACE} -l llm-d.ai/role=prefill
kubectl get pods -n ${NAMESPACE} -l llm-d.ai/role=decode

11. 验证投放

在路由器、预填充和解码工作负载运行的情况下，您现在可以验证服务 API。

运行验证脚本：

./11_verify_serving.sh

工作原理：

  • 该脚本会检索 GKE 网关的内部 IP。
  • 它会启动一个临时 pod (curl-debug-comp)，以向 http://${GATEWAY_IP}/v1/completions 发送完成请求。
  • 它会启动另一个 pod (curl-debug-chat) 向 http://${GATEWAY_IP}/v1/chat/completions 发送聊天请求。

您应该会看到来自 Qwen 模型的 JSON 成功响应：

{
  "choices": [
    {
      "text": "... [Model Response] ..."
    }
  ]
}

12. 清理

为避免系统向您的 Google Cloud 账号持续收取费用，请删除在此 Codelab 中创建的资源。

运行收尾清理脚本：

./12_teardown_cleanup.sh

此脚本的作用：

  1. 删除 GKE 节点池（tpu7-pool-1、tpu7-pool-2）。
  2. 删除 GKE 集群 (tpu-serving-cluster)。
  3. 删除资源政策 (superslice-policy)。
  4. 删除 VPC 网络 (qwen-serving-main)。

或者，如果您为此 Codelab 创建了专用项目，则可以删除整个项目：

gcloud projects delete ${PROJECT_ID}

13. 恭喜

恭喜！您已成功探索 GKE 动态切分，并部署了分离式 LLM 服务架构。

您学到的内容

  • 如何启用 GKE Slice Controller 并为增量配置配置节点池。
  • 如何使用 Kueue 请求特定的 TPU 拓扑。
  • 子切片如何拆分大型 TPU 块以用于较小、独立的 JAX 工作负载。
  • 超级切片如何将多个节点池拼接成一个更大的虚拟 TPU 切片。
  • 如何使用 LWS、Gateway API 和 vLLM 部署预填充/解码分离式服务。

参考文档