GKE 中的 TPU 简介

Google Kubernetes Engine (GKE) 客户现在可以创建包含 TPU v4 和 v5e Pod 的 Kubernetes 节点池。TPU Pod 是通过高速互连连接的一组 TPU 设备。对于不需要完整 TPU Pod 的工作负载，您可以使用完整 TPU Pod 的子集，称为 TPU 切片。与完整的 TPU Pod 一样，切片中的每个 TPU 设备都有自己的 TPU 虚拟机。我们将 TPU 虚拟机及其连接的设备称为主机或 TPU 节点。如需详细了解 TPU Pod，请参阅系统架构。

由于 GKE 上下文中使用的术语 Pod 通常表示 Kubernetes Pod，为避免混淆，我们始终将一个或多个 TPU 设备的集合称为“切片”。

使用 GKE 时，您必须先创建 GKE 集群。

然后，向集群添加节点池。GKE 节点池是具有相同特性的虚拟机集合。对于 TPU 工作负载，节点池由 TPU 虚拟机组成。

节点池类型

GKE 支持两种类型的 TPU 节点池：

多主机 TPU 切片节点池

多主机 TPU 切片节点池是包含两个或更多互连 TPU 虚拟机的节点池。每个虚拟机都有一个与之连接的 TPU 设备。多主机切片中的 TPU 通过高速互连 (ICI) 连接。多主机 TPU 切片节点池不可变。创建多主机切片节点池后，您便无法向其添加节点。例如，您无法在创建 v4-32 节点池后再向该节点池添加额外的 Kubernetes 节点 (TPU 虚拟机)。如需向 GKE 集群添加其他 TPU 切片，您必须创建新的节点池。

多主机 TPU 切片节点池中的主机被视为单个原子单元。如果 GKE 无法在切片中部署一个节点，则切片中的所有节点都将无法部署。

如果需要修复多主机 TPU 切片中的节点，GKE 将关停切片中的所有 TPU 虚拟机（节点），强制逐出工作负载中的所有 GKE Pod。一旦切片中的所有 TPU 虚拟机都已启动并运行，便可在新切片中的 TPU 虚拟机上调度 GKE Pod。

下图显示了一个 TPU v5litepod-16 (v5e) 多主机切片示例。此切片有四个 TPU 虚拟机。每个 TPU 虚拟机有四个 TPU v5e 芯片，通过高速互连 (ICI) 连接，每个 TPU v5e 芯片有一个 TensorCore。

多主机 TPU 切片图

下图展示了一个 GKE 集群，其中包含一个 TPU v5litepod-16 (v5e) 切片（拓扑：4x4）和一个 TPU v5litepod-8 (v5e) 切片（拓扑：2x4）：

TPU v5e Pod 示意图

如需查看在多主机 TPU 切片上运行工作负载的示例，请参阅在多主机 TPU 切片上运行工作负载。

单主机 TPU 切片节点池

单主机切片节点池是包含一个或多个独立 TPU 虚拟机的节点池。其中的每个虚拟机都有一个与之连接的 TPU 设备。虽然单主机切片节点池中的虚拟机可以通过数据中心网络 (DCN) 进行通信，但挂接到虚拟机的 TPU 不互连。下图显示了一个包含七个 v4-8 机器的单主机 TPU 切片的示例：

单主机切片节点池示意图

如需查看在单主机 TPU 切片上运行工作负载的示例，请参阅在 TPU 节点上运行工作负载。

适用于 GKE 节点池的 TPU 机器类型

在创建节点池之前，您需要选择工作负载所需的 TPU 版本和 TPU 切片的大小。GKE Standard 版本 1.26.1-gke.1500 及更高版本支持 TPU v4，GKE Standard 版本 1.27.2-gke.2100 及更高版本支持 v5e，GKE Standard 版本 1.28.3-gke.1024000 及更高版本支持 v5p。

GKE Autopilot 版本 1.29.2-gke.1521000 及更高版本支持 TPU v4、v5e 和 v5p。

如需详细了解不同 TPU 版本的硬件规格，请参阅系统架构。创建 TPU 节点池时，请根据模型的大小及其所需的内存选择 TPU 切片大小（TPU 拓扑）。您在创建节点池时指定的机器类型取决于切片的版本和大小。

v5e

以下是训练和推断用例支持的 TPU v5e 机器类型和拓扑：

机器类型 拓扑 TPU 芯片的数量 虚拟机数量 推荐用例
ct5lp-hightpu-1t 1x1 1 1 训练、单主机推理
ct5lp-hightpu-4t 2x2 4 1 训练、单主机推理
ct5lp-hightpu-8t 2x4 8 1 训练、单主机推理
ct5lp-hightpu-4t 2x4 8 2 训练、多主机推理
ct5lp-hightpu-4t 4x4 16 4 大规模训练、多主机推理
ct5lp-hightpu-4t 4x8 32 8 大规模训练、多主机推理
ct5lp-hightpu-4t 8x8 64 16 大规模训练、多主机推理
ct5lp-hightpu-4t 8x16 128 32 大规模训练、多主机推理
ct5lp-hightpu-4t 16x16 256 64 大规模训练、多主机推理

Cloud TPU v5e 是一款训练和推理相结合的产品。训练作业针对吞吐量和可用性进行了优化，而推理作业针对延迟时间进行了优化。如需了解详情，请参阅 v5e 训练加速器类型和 v5e 推断加速器类型。

TPU v5e 机器提供 us-west4-a、us-east5-b 和 us-east1-c 版本。GKE Standard 集群必须运行控制平面 1.27.2-gke.2100 或更高版本。GKE Autopilot 必须运行控制平面 1.29.2-gke.1521000 或更高版本。如需详细了解 v5e，请参阅 Cloud TPU v5e 训练。

机器类型比较：

机器类型 ct5lp-hightpu-1t ct5lp-hightpu-4t ct5lp-hightpu-8t
v5e 芯片数量 1 4 8
vCPU 的数量 24 112 224
RAM (GB) 48 192 384
NUMA 节点的数量 1 1 2
抢占可能性 高 中 低

为了给具有更多芯片的虚拟机腾出空间，GKE 调度器可能会抢占和重新安排具有较少芯片的虚拟机。因此，8 芯片虚拟机可能会抢占 1 芯片和 4 芯片虚拟机。

v4 和 v5p

以下是 TPU v4 和 v5p 机器类型：

机器类型 vCPU 的数量 内存 (GB) NUMA 节点的数量
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

创建 TPU v4 切片时，请使用 ct4p-hightpu-4t 机器类型，它具有一个主机和包含 4 个芯片。如需了解详情，请参阅 v4 拓扑和 TPU 系统架构。us-central2-b 提供 TPU v4 Pod 机器类型。您的 GKE Standard 集群必须运行 1.26.1-gke.1500 或更高版本的控制平面。GKE Autopilot 集群必须运行 1.29.2-gke.1521000 或更高版本的控制平面。

创建 TPU v5p 切片时，请使用 ct5p-hightpu-4t 机器类型，它具有一个主机和包含 4 个芯片。us-west4-a 和 us-east5-a 提供 TPU v5p Pod 机器类型。GKE Standard 集群必须运行控制平面 1.28.3-gke.1024000 或更高版本。GKE Autopilot 必须运行 1.29.2-gke.1521000 或更高版本。如需详细了解 v5p，请参阅 v5p 培训简介。

已知问题和限制

  • Docker pod 数量上限：单个 TPU 虚拟机中最多可以运行 256 个 Docker pod。
  • 仅限特定预留：在 GKE 中使用 TPU 时，SPECIFIC 是 gcloud container node-pools create 命令的 --reservation-affinity 标志唯一支持的值。
  • 仅支持抢占式 TPU 的 Spot 虚拟机变体：Spot 虚拟机与抢占式虚拟机类似，受到相同的可用性限制，但没有 24 小时的最长持续时间。
  • 不支持费用分配：GKE 费用分配和用量计量不包含任何与 TPU 用量或费用相关的数据。
  • 自动扩缩器可能会计算容量：在新的 TPU 节点可用之前，集群自动扩缩器可能会错误地计算这些节点的容量。然后，集群自动扩缩器可能会执行额外的纵向扩容，因此产生的节点数量会超出需求。在执行常规缩减操作后，集群自动扩缩器将缩减其他不需要的节点。
  • 自动伸缩器取消纵向扩容：集群自动伸缩器取消对处于等待状态超过 15 分钟的 TPU 节点池的纵向扩容。集群自动扩缩器将稍后重试此类纵向扩容操作。对于不使用预留的客户，此行为可能会降低 TPU 的可用性。
  • 污点可能阻止纵向缩容：对 TPU 污点有容忍的非 TPU 工作负载可能会在 TPU 节点池排空期间重新创建节点池，导致这些工作负载无法纵向缩容。

确保有足够的 TPU 和 GKE 配额

您可能需要在创建资源的区域中增加某些与 GKE 相关的配额。

以下配额具有默认值，可能需要增加：

  • Persistent Disk SSD (GB) 配额：默认情况下，每个 Kubernetes 节点的启动磁盘需要 100GB。因此，此配额应至少设置为（预计创建的 GKE 节点数量上限）* 100GB。
  • 已用 IP 地址配额：每个 Kubernetes 节点使用一个 IP 地址。因此，此配额应至少设置为您预计创建的 GKE 节点数上限。

如需申请增加配额，请参阅申请更高配额。如需详细了解 TPU 配额类型，请参阅 TPU 配额。

您的配额增加申请可能需要几天时间才能获得批准。如果配额增加请求在几天内获批时遇到任何问题，请与您的 Google 客户支持团队联系。

迁移 TPU 预留

如果您不打算在 GKE 中将现有 TPU 预留与 TPU 搭配使用，请跳过此部分并转到创建 Google Kubernetes Engine 集群。

为了将预留的 TPU 与 GKE 搭配使用，您必须先将 TPU 预留迁移到基于 Compute Engine 的新预留系统。

关于此次迁移，有一些重要注意事项：

  • 迁移到基于 Compute Engine 的新预留系统的 TPU 容量不能与 Cloud TPU Queued Resource API 搭配使用。如果您打算在预留中使用 TPU 排队的资源，则只需要将部分 TPU 预留迁移到基于 Compute Engine 的新预留系统。
  • 将 TPU 迁移到基于 Compute Engine 的新预留系统时，任何工作负载都无法活跃在 TPU 上。
  • 选择执行迁移的时间，并与您的 Google Cloud 客户支持团队合作安排迁移。迁移时间范围应为工作时间（周一至周五，上午 9 点至下午 5 点（太平洋时间）。

创建一个 Google Kubernetes Engine 集群

请参阅 Google Kubernetes Engine 文档中的创建集群。

创建 TPU 节点池

请参阅 Google Kubernetes Engine 文档中的创建节点池。

在没有特权模式的情况下运行

如果要缩小容器的权限范围，请参阅 TPU 权限模式。

在 TPU 节点上运行工作负载

请参阅 Google Kubernetes Engine 文档中的在 TPU 节点上运行工作负载。

节点选择器

为了让 Kubernetes 在 TPU 节点上调度您的工作负载，您必须为 Google Kubernetes Engine 清单中的每个 TPU 节点指定两个选择器：

  • 将 cloud.google.com/gke-accelerator-type 设置为 tpu-v5-lite-podslice 或 tpu-v4-podslice。
  • 将 cloud.google.com/gke-tpu-topology 设置为 TPU 节点的 TPU 拓扑。

训练工作负载和推断工作负载部分包含示例清单，说明了如何使用这些节点选择器。

工作负载安排注意事项

TPU 具有独特的特征，这就要求在 Kubernetes 中进行特殊的工作负载调度和管理。如需了解详情，请参阅 GKE 文档中的工作负载调度注意事项。

TPU 节点修复

如果多主机 TPU 切片节点池中的 TPU 节点健康状况不佳，则系统会重新创建整个节点池。如需了解详情，请参阅 GKE 文档中的节点自动修复。

多切片 - 不局限于单个切片

您可以将较小的切片汇总到一个多切片中，以处理较大的训练工作负载。如需了解详情，请参阅 Cloud TPU 多切片。

训练工作负载教程

这些教程重点介绍了如何在多主机 TPU 切片（例如，4 台 v5e 机器）上训练工作负载。它们涵盖以下模型：

  • 拥抱 FLAX 模型：《Pokémon》中的 Train Diffusion
  • PyTorch/XLA：WikiText 上的 GPT2

下载教程资源

使用以下命令下载每个预训练模型的教程 Python 脚本和 YAML 规范：

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

创建并连接到集群

创建区域级 GKE 集群，以便将 Kubernetes 控制平面复制到三个可用区，从而提供更高的可用性。根据您使用的 TPU 版本，在 us-west4、us-east1 或 us-central2 中创建集群。如需详细了解 TPU 和地区，请参阅 Cloud TPU 区域和地区。

以下命令会创建一个新的 GKE 区域级集群，该集群已订阅快速发布渠道，其节点池最初在每个可用区包含一个节点。该命令还在集群上启用 Workload Identity 和 Cloud Storage FUSE CSI 驱动程序功能，因为本指南中的示例推断工作负载使用 Cloud Storage 存储分区来存储预训练模型。

gcloud container clusters create cluster-name \
  --region your-region \
  --release-channel rapid \
  --num-nodes=1 \
  --workload-pool=project-id.svc.id.goog \
  --addons GcsFuseCsiDriver

如需为现有集群启用 Workload Identity 和 Cloud Storage FUSE CSI 驱动程序功能，请运行以下命令：

gcloud container clusters update cluster-name \
  --region your-region \
  --update-addons GcsFuseCsiDriver=ENABLED \
  --workload-pool=project-id.svc.id.goog

示例工作负载配置了以下假设：

  • 节点池正在使用“tpu-topology=4x4”，有四个节点
  • 节点池正在使用machine-type ct5lp-hightpu-4t

运行以下命令以连接到新创建的集群：

gcloud container clusters get-credentials cluster-name \
--location=cluster-region

拥抱 FLAX 模型：《Pokémon》中的 Train Diffusion

此示例使用 Pokémon 数据集通过 HuggingFace 训练 Stable Diffusion 模型。

稳定扩散模型是一种潜在的文本到图像模型，可根据任何文本输入生成逼真的图像。如需详细了解 Stable Diffusion，请参阅：

创建 Docker 映像

Dockerfile 位于文件夹 ai-on-gke/gke-tpu-examples/training/diffusion/ 下。运行以下命令来构建和推送 Docker 映像。

cd ai-on-gke/gke-tpu-examples/training/diffusion/
docker build -t gcr.io/project-id/diffusion:latest .
docker push gcr.io/project-id/diffusion:latest

部署工作负载

创建一个包含以下内容的文件，并将其命名为 tpu_job_diffusion.yaml。 使用您刚刚创建的图片填写图片字段。

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-diffusion
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-diffusion
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (e.g. 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-diffusion
        image: gcr.io/<var>project-name</var>/diffusion:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU usage metrics, if supported
        command:
        - bash
        - -c
        - |
          cd examples/text_to_image
          python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

然后使用以下命令进行部署：

kubectl apply -f tpu_job_diffusion.yaml

清理

作业运行完毕后，您可以使用以下命令将其删除：

kubectl delete -f tpu_job_diffusion.yaml

PyTorch/XLA：WikiText 上的 GPT2

本教程介绍如何使用 Wikitext 数据集在 PyTorch/XLA 上使用 HuggingFace 在 v5e TPU 上运行 GPT2。

创建 Docker 映像

Dockerfile 位于文件夹 ai-on-gke/gke-tpu-examples/training/gpt/ 下。运行以下命令来构建和推送 Docker 映像。

cd ai-on-gke/gke-tpu-examples/training/gpt/
docker build -t gcr.io/project-id/gpt:latest .
docker push gcr.io/project-id/gpt:latest

部署工作负载

复制以下 YAML 并将其保存在名为 tpu_job_gpt.yaml 的文件中。在映像字段中填写您刚刚创建的映像。

kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-gpt
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-gpt
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (for example, 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-gpt
        image: gcr.io/<var>project-name</var>/gpt:latest
        ports:
        - containerPort: 8479
        - containerPort: 8478
        - containerPort: 8477
        - containerPort: 8476
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        env:
        - name: PJRT_DEVICE
          value: 'TPU_C_API'
        - name: XLA_USE_BF16
          value: '1'
        command:
        - bash
        - -c
        - |
          numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py   --num_cores 4 examples/pytorch/language-modeling/run_clm.py    --num_train_epochs 3 --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 16    --per_device_eval_batch_size 16 --do_train --do_eval  --output_dir /tmp/test-clm     --overwrite_output_dir --config_name my_config_2.json --cache_dir /tmp --tokenizer_name gpt2  --block_size 1024 --optim adafactor --adafactor true --save_strategy no --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

使用以下命令部署工作流：

kubectl apply -f tpu_job_gpt.yaml

清理

作业运行完毕后，您可以使用以下命令将其删除：

kubectl delete -f tpu_job_gpt.yaml

教程：单主机推理工作负载

本教程介绍如何使用 JAX、TensorFlow 和 PyTorch 在 GKE v5e TPU 上为预训练模型运行单主机推理工作负载。概括来讲，需要在 GKE 集群上执行四个单独的步骤：

  1. 创建一个 Cloud Storage 存储桶并设置对该存储桶的访问权限。您可以使用 Cloud Storage 存储桶来存储预训练模型。

  2. 下载预训练模型并将其转换为与 TPU 兼容的模型。应用 GKE Pod，该 Pod 可以下载预训练模型、使用 Cloud TPU Converter 以及使用 Cloud Storage FUSE CSI 驱动程序将转换后的模型存储到 Cloud Storage 存储桶中。Cloud TPU Converter 不需要专门的硬件。本教程介绍如何下载模型并在 CPU 节点池中运行 Cloud TPU 转换器。

  3. 启动转换后的模型的服务器。应用一个部署，该部署使用服务器框架（由存储在 ReadOnly Many (ROX) 永久性卷中的卷提供支持）来提供模型。部署副本必须在 v5e Pod TPU 节点上运行，每个节点有一个 Kubernetes Pod。

  4. 部署负载均衡器以测试模型服务器。使用 LoadBalancer 服务向外部请求公开服务器。我们提供了 Python 脚本，其中包含用于测试模型服务器的示例请求。

下图展示了负载平衡器如何路由请求。

显示负载均衡器路由的图表

服务器部署示例

这些示例工作负载配置了以下假设：

  • 集群正在使用包含 3 个节点的 TPU v5 节点池运行
  • 节点池使用的是 ct5lp-hightpu-1t 机器类型，其中：
    • 拓扑为 1x1
    • TPU 芯片的数量为 1

以下 GKE 清单定义了单个主机服务器 Deployment。

apiVersion: apps/v1
kind: Deployment
metadata:
  name: bert-deployment
spec:
  selector:
    matchLabels:
      app: tf-bert-server
  replicas: 3 # number of nodes in node pool
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
      labels:
        app: tf-bert-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1  # target topology
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice  # target version
      containers:
      - name: serve-bert
        image: us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
        env:
        - name: MODEL_NAME
          value: "bert"
        volumeMounts:
        - mountPath: "/models/"
          name: bert-external-storage
        ports:
        - containerPort: 8500
        - containerPort: 8501
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        resources:
          requests:
            google.com/tpu: 1 # TPU chip request
          limits:
            google.com/tpu: 1 # TPU chip request
      volumes:
      - name: bert-external-storage
        persistentVolumeClaim:
          claimName: external-storage-pvc

如果您在 TPU 节点池中使用的节点数量不同，请将 replicas 字段更改为节点数。

如果您的 Standard 集群运行 GKE 1.27 或更低版本，请将以下字段添加到清单中：

spec:
  securityContext:
    privileged: true

在 GKE 1.28 或更高版本中，您无需在特权模式下运行 Pod。如需了解详情，请参阅在没有特权模式的情况下运行容器。

如果您使用的是其他机器类型，请执行以下操作：

  • 将 cloud.google.com/gke-tpu-topology 设置为您所用机器类型的拓扑。
  • 设置 resources 下的两个 google.com/tpu 字段，以匹配相应机器类型的芯片数量。

设置

使用以下命令下载教程中的 Python 脚本和 YAML 清单：

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

转到 single-host-inference 目录：

cd ai-on-gke/gke-tpu-examples/single-host-inference/

设置 Python 环境

您在本教程中使用的 Python 脚本需要 Python 3.9 或更高版本。在运行 Python 测试脚本之前，请务必为每个教程安装 requirements.txt。

如果您没有在本地环境中进行正确的 Python 设置，则可以使用 Cloud Shell 下载并运行本教程中的 Python 脚本。

设置集群

  1. 使用 e2-standard-4 机器类型创建集群。

    gcloud container clusters create cluster-name \
--region your-region \
--release-channel rapid \
--num-nodes=1 \
--machine-type=e2-standard-4 \
--workload-pool=project-id.svc.id.goog \
--addons GcsFuseCsiDriver

  2. 创建单主机 TPU 节点池。

示例工作负载做出如下假设：

  • 您的集群正在运行包含 3 个节点的 TPU v5e 节点池。
  • TPU 节点池正在使用机器类型 ct5lp-hightpu-1t。

如果您使用的集群配置与上述不同，则需要修改服务器部署清单。

对于 JAX Stable Diffusion 演示，您需要一个具有 16 Gi+ 可用内存的机器类型的 CPU 节点池（例如 e2-standard-4）。可通过 gcloud container clusters create 命令或使用以下命令向现有集群添加额外的节点池进行配置：

gcloud beta container node-pools create your-pool-name \
  --zone=your-cluster-zone \
  --cluster=your-cluster-name \
  --machine-type=e2-standard-4 \
  --num-nodes=1

替换以下内容：

  • your-pool-name：要创建的节点池的名称。
  • your-cluster-zone：创建集群的可用区。
  • your-cluster-name：要在其中添加节点池的集群的名称。
  • your-machine-type：要在节点池中创建的节点的机器类型。

设置模型存储

您可以通过多种方式存储模型以提供服务。在本教程中，我们将使用以下方法：

  • 为了将预训练模型转换为在 TPU 上运行，我们将使用由具有 ReadWriteMany (RWX) 访问权限的 Persistent Disk 提供支持的虚拟私有云。
  • 为了在多个单主机 TPU 上提供模型，我们将使用 Cloud Storage 存储桶支持的同一 VPC。

运行以下命令以创建 Cloud Storage 存储桶。

gcloud storage buckets create gs://your-bucket-name \
  --project=your-bucket-project-id \
  --location=your-bucket-location

替换以下内容：

  • your-bucket-name：Cloud Storage 存储桶的名称。
  • your-bucket-project-id：您在其中创建 Cloud Storage 存储桶的项目的 ID。
  • your-bucket-location：Cloud Storage 存储桶的位置。如需提高性能，请指定 GKE 集群的运行位置。

请按照以下步骤向 GKE 集群授予对存储桶的访问权限。为了简化设置，以下示例使用默认命名空间和默认 Kubernetes 服务帐号。如需了解详情，请参阅使用 GKE Workload Identity 配置对 Cloud Storage 存储分区的访问权限。

  1. 为您的应用创建 IAM 服务账号，或使用现有 IAM 服务账号。您可以在 Cloud Storage 存储桶的项目中使用任何 IAM 服务帐号。

    gcloud iam service-accounts create your-iam-service-acct \
--project=your-bucket-project-id

    替换以下内容：

    • your-iam-service-acct：新 IAM 服务帐号的名称。
    • your-bucket-project-id：您在其中创建了 IAM 服务帐号的项目的 ID。IAM 服务帐号必须与 Cloud Storage 存储桶位于同一项目中。

  2. 确保您的 IAM 服务帐号具有您所需的存储角色。

    gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
--member "serviceAccount:your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com" \
--role "roles/storage.objectAdmin"

    替换以下内容：

    • your-bucket-name - Cloud Storage 存储桶的名称。
    • your-iam-service-acct：新 IAM 服务帐号的名称。
    • your-bucket-project-id：您在其中创建了 IAM 服务帐号的项目的 ID。

  3. 通过在两个服务帐号之间添加 IAM 政策绑定，允许 Kubernetes 服务帐号模拟 IAM 服务帐号。此绑定允许 Kubernetes 服务帐号充当 IAM 服务帐号。

    gcloud iam service-accounts add-iam-policy-binding your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:your-project-id.svc.id.goog[default/default]"

    替换以下内容：

    • your-iam-service-acct：新 IAM 服务帐号的名称。
    • your-bucket-project-id：您在其中创建了 IAM 服务帐号的项目的 ID。
    • your-project-id：您在其中创建了 GKE 集群的项目的 ID。您的 Cloud Storage 存储分区和 GKE 集群可以位于同一项目中，也可以位于不同的项目中。

  4. 使用 IAM 服务账号的电子邮件地址为 Kubernetes 服务账号添加注解。

    kubectl annotate serviceaccount default \
  --namespace default \
  iam.gke.io/gcp-service-account=your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com

    替换以下内容：

    • your-iam-service-acct：新 IAM 服务帐号的名称。
    • your-bucket-project-id：您在其中创建了 IAM 服务帐号的项目的 ID。

  5. 运行以下命令，在此演示的 YAML 文件中填充您的存储桶名称：

    find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"

    将 your-bucket-name 替换为 Cloud Storage 存储桶的名称。

  6. 使用以下命令创建永久性卷和永久性卷声明：

    kubectl apply -f pvc-pv.yaml

JAX 模型推断与服务

安装 Python 依赖项，以便运行向 JAX 模型服务发送请求的教程 Python 脚本。

pip install -r jax/requirements.txt

运行 JAX BERT E2E 服务演示：

此演示使用的是 Hugging Face 中的预训练 BERT 模型。

Kubernetes Pod 执行以下步骤：

  1. 从示例资源下载 Python 脚本 export_bert_model.py，然后使用该脚本将预训练的 bert 模型下载到临时目录。
  2. 使用 Cloud TPU Converter 图片将预训练模型从 CPU 转换为 TPU，并将模型存储在您在setup期间创建的 Cloud Storage 存储桶中。

此 Kubernetes Pod 配置为在默认节点池 CPU 上运行。使用以下命令运行 Pod：

kubectl apply -f jax/bert/install-bert.yaml

使用以下命令验证模型是否已正确安装：

kubectl get pods install-bert

STATUS 可能需要几分钟时间才能读取到 Completed。

启动模型的 TF 模型服务器

本教程中的示例工作负载假定：

  • 集群正在使用包含三个节点的 TPU v5 节点池运行
  • 节点池使用的 ct5lp-hightpu-1t 机器类型包含一个 TPU 芯片。

如果您使用的集群配置与上述不同，则需要修改服务器部署清单。

应用部署
kubectl apply -f jax/bert/serve-bert.yaml

使用以下命令验证服务器是否正在运行：

kubectl get deployment bert-deployment

AVAILABLE 可能需要 1 分钟的时间才能读出 3。

应用负载均衡器服务
kubectl apply -f jax/bert/loadbalancer.yaml

使用以下命令验证负载均衡器是否已准备好接收外部流量：

kubectl get svc tf-bert-service

EXTERNAL_IP 可能需要几分钟时间才能列出 IP。

向模型服务器发送请求

从负载均衡器服务获取外部 IP：

EXTERNAL_IP=$(kubectl get services tf-bert-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

运行用于向服务器发送请求的脚本：

python3 jax/bert/bert_request.py $EXTERNAL_IP

预期输出：

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
清理

如需清理资源，请以倒序运行 kubectl delete。

kubectl delete -f jax/bert/loadbalancer.yaml
kubectl delete -f jax/bert/serve-bert.yaml
kubectl delete -f jax/bert/install-bert.yaml

运行 JAX Stable Diffusion E2E 服务演示

本演示使用 Hugging Face 中的预训练稳定扩散模型。

从 Flax Stable Diffusion 模型中导出与 TPU 兼容的 TF2 保存模型

导出稳定扩散模型需要集群的 CPU 节点池具有具有 16Gi+ 可用内存的机器类型（如设置集群中所述）。

Kubernetes Pod 执行以下步骤：

  1. 从示例资源中下载 Python 脚本 export_stable_diffusion_model.py，然后使用该脚本将预训练的稳定扩散模型下载到临时目录。
  2. 使用 Cloud TPU Converter 图片将预训练模型从 CPU 转换为 TPU，并将模型存储在您在设置存储空间期间创建的 Cloud Storage 存储桶中。

此 Kubernetes Pod 配置为在默认 CPU 节点池池上运行。使用以下命令运行 Pod：

kubectl apply -f jax/stable-diffusion/install-stable-diffusion.yaml

使用以下命令验证模型是否已正确安装：

kubectl get pods install-stable-diffusion

STATUS 可能需要几分钟时间才能读取到 Completed。

启动模型的 TF 模型服务器容器

示例工作负载已使用以下假设进行配置：

  • 集群在使用包含三个节点的 TPU v5 节点池运行
  • 节点池使用的是 ct5lp-hightpu-1t 机器类型，其中：
    • 拓扑为 1x1
    • TPU 芯片的数量为 1

如果您使用的集群配置与上述不同，则需要修改服务器部署清单。

应用部署：

kubectl apply -f jax/stable-diffusion/serve-stable-diffusion.yaml

验证服务器是否按预期运行：

kubectl get deployment stable-diffusion-deployment

AVAILABLE 可能需要 1 分钟的时间才能读出 3。

应用负载均衡器服务：

kubectl apply -f jax/stable-diffusion/loadbalancer.yaml

使用以下命令验证负载均衡器是否已准备好接收外部流量：

kubectl get svc tf-stable-diffusion-service

EXTERNAL_IP 可能需要几分钟时间才能列出 IP。

向模型服务器发送请求

从负载均衡器获取外部 IP：

EXTERNAL_IP=$(kubectl get services tf-stable-diffusion-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

运行用于向服务器发送请求的脚本

python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP

预期输出：

提示为 Painting of a squirrel skating in New York，输出图片将保存到当前目录中的 stable_diffusion_images.jpg。

清理

如需清理资源，请以倒序运行 kubectl delete。

kubectl delete -f jax/stable-diffusion/loadbalancer.yaml
kubectl delete -f jax/stable-diffusion/serve-stable-diffusion.yaml
kubectl delete -f jax/stable-diffusion/install-stable-diffusion.yaml

运行 TensorFlow ResNet-50 E2E 服务演示：

安装 Python 依赖项，以便运行向 TF 模型服务发送请求的教程 Python 脚本。

pip install -r tf/resnet50/requirements.txt
第 1 步：转换模型

应用模型转换：

kubectl apply -f tf/resnet50/model-conversion.yml

使用以下命令验证模型是否已正确安装：

kubectl get pods resnet-model-conversion

STATUS 可能需要几分钟时间才能读取到 Completed。

第 2 步：使用 TensorFlow Serving 提供模型

应用模型服务部署：

kubectl apply -f tf/resnet50/deployment.yml

使用以下命令验证服务器是否按预期运行：

kubectl get deployment resnet-deployment

AVAILABLE 可能需要 1 分钟的时间才能读出 3。

应用负载均衡器服务：

kubectl apply -f tf/resnet50/loadbalancer.yml

使用以下命令验证负载均衡器是否已准备好接收外部流量：

kubectl get svc resnet-service

EXTERNAL_IP 可能需要几分钟时间才能列出 IP。

第 3 步：向模型服务器发送测试请求

从负载均衡器获取外部 IP：

EXTERNAL_IP=$(kubectl get services resnet-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

运行测试请求 (HTTP) 脚本以向模型服务器发送请求。

python3 tf/resnet50/request.py --host $EXTERNAL_IP

响应应如下所示：

Predict result: ['ImageNet ID: n07753592, Label: banana, Confidence: 0.94921875',
'ImageNet ID: n03532672, Label: hook, Confidence: 0.0223388672', 'ImageNet ID: n07749582,
Label: lemon, Confidence: 0.00512695312
第 4 步：清理

如需清理资源，请运行以下 kubectl delete 命令：

kubectl delete -f tf/resnet50/loadbalancer.yml
kubectl delete -f tf/resnet50/deployment.yml
kubectl delete -f tf/resnet50/model-conversion.yml

请务必在使用完 GKE 节点池和集群后将其删除。

PyTorch 模型推断和服务

安装 Python 依赖项，以便运行向 PyTorch 模型服务发送请求的教程 Python 脚本：

pip install -r pt/densenet161/requirements.txt

运行 TorchServe Densenet161 E2E 服务演示：

  1. 生成模型归档文件。

    1. 应用模型归档：
    kubectl apply -f pt/densenet161/model-archive.yml
    1. 使用以下命令验证模型是否已正确安装：
    kubectl get pods densenet161-model-archive

    STATUS 可能需要几分钟时间才能读取到 Completed。

  2. 使用 TorchServe 模型：

    1. 应用模型服务部署：

      kubectl apply -f pt/densenet161/deployment.yml

    2. 使用以下命令验证服务器是否按预期运行：

      kubectl get deployment densenet161-deployment

      AVAILABLE 可能需要 1 分钟的时间才能读出 3。

    3. 应用负载均衡器服务：

      kubectl apply -f pt/densenet161/loadbalancer.yml

      使用以下命令验证负载均衡器是否已准备好接收外部流量：

      kubectl get svc densenet161-service

      EXTERNAL_IP 可能需要几分钟时间才能列出 IP。

  3. 向模型服务器发送测试请求：

    1. 从负载均衡器获取外部 IP：

      EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

    2. 运行测试请求脚本，向模型服务器发送请求 (HTTP)：

      python3 pt/densenet161/request.py --host $EXTERNAL_IP

      您应该会看到类似如下的响应：

      Request successful. Response: {'tabby': 0.47878125309944153, 'lynx': 0.20393909513950348, 'tiger_cat': 0.16572578251361847, 'tiger': 0.061157409101724625, 'Egyptian_cat': 0.04997897148132324

  4. 通过运行以下 kubectl delete 命令清理资源：

    kubectl delete -f pt/densenet161/loadbalancer.yml
kubectl delete -f pt/densenet161/deployment.yml
kubectl delete -f pt/densenet161/model-archive.yml

    请务必在使用完 GKE 节点池和集群后将其删除。

排查常见问题

您可以在排查 GKE 中的 TPU 问题中找到 GKE 问题排查信息。

TPU 初始化失败

如果您遇到以下错误，请确保在特权模式下运行 TPU 容器，或者已在容器内增加了 ulimit。如需了解详情，请参阅在没有特权模式的情况下运行。

TPU platform initialization failed: FAILED_PRECONDITION: Couldn't mmap: Resource
temporarily unavailable.; Unable to create Node RegisterInterface for node 0,
config: device_path:      "/dev/accel0" mode: KERNEL debug_data_directory: ""
dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true;
could not create driver instance

调度出现死锁

假设您有两个作业（作业 A 和作业 B），它们都将安排在具有给定 TPU 拓扑（例如 v4-32）的 TPU 切片上。另外，假设您在 GKE 集群中有两个 v4-32 TPU 切片；我们将这两个作业称为切片 X 和切片 Y。由于您的集群有足够的容量来调度这两个作业，因此理论上两个作业都应快速调度 - 两个 TPU v4-32 切片各有一个作业。

但是，如果不仔细规划，可能会导致调度死锁。假设 Kubernetes 调度器在切片 X 上调度作业 A 中的一个 Pod，然后在切片 X 上调度作业 B 中的一个 Pod。在这种情况下，根据作业 A 的 Pod 亲和性规则，调度器将尝试在切片 X 上为作业 A 调度所有剩余 Pod。对于作业 B 也是如此。因此，作业 A 和作业 B 都无法在单个切片上完全调度。这会导致调度死锁。

为了避免发生调度死锁的风险，您可以使用 Pod 反亲和性，并将 cloud.google.com/gke-nodepool 作为 topologyKey，如以下示例所示：

apiVersion: batch/v1
kind: Job
metadata:
 name: pi
spec:
 parallelism: 2
 template:
   metadata:
     labels:
       job: pi
   spec:
     affinity:
       podAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: In
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
       podAntiAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: NotIn
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
           namespaceSelector:
             matchExpressions:
             - key: kubernetes.io/metadata.name
               operator: NotIn
               values:
               - kube-system
     containers:
     - name: pi
       image: perl:5.34.0
       command: ["sleep",  "60"]
     restartPolicy: Never
 backoffLimit: 4

使用 Terraform 创建 TPU 节点池资源

您还可以使用 Terraform 管理集群和节点池资源。

在现有 GKE 集群中创建多主机 TPU 切片节点池

如果您在现有集群中创建多主机 TPU 节点池，则可以使用以下 Terraform 代码段：

resource "google_container_cluster" "cluster_multi_host" {
  …
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
    workload_pool = "my-gke-project.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "multi_host_tpu" {
  provider           = google-beta
  project            = "<var>your-project</var>"
  name               = "<var>your-node-pool</var>"
  location           = "<var>us-central2</var>"
  node_locations     = ["<var>us-central2-b</var>"]
  cluster            = google_container_cluster.cluster_multi_host.name
  initial_node_count = 2

  node_config {
    machine_type = "ct4p-hightpu-4t"
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = ["<var>your-reservation-name</var>"]
    }
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = "2x2x2"
  }
}

替换以下值：

  • your-project：在其中运行工作负载的 Google Cloud 项目。
  • your-node-pool：您要创建的节点池的名称。
  • us-central2：运行工作负载的区域。
  • us-central2-b：运行工作负载的可用区。
  • your-reservation-name：预留的名称。

在现有 GKE 集群中创建单主机 TPU 切片节点池

使用以下 Terraform 代码段：

resource "google_container_cluster" "cluster_single_host" {
  …
  cluster_autoscaling {
    autoscaling_profile = "OPTIMIZE_UTILIZATION"
  }
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
  workload_pool = "<var>your-project-id</var>.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "single_host_tpu" {
  provider           = google-beta
  project            = "<var>your-project</var>"
  name               = "<var>your-node-pool</var>"
  location           = "<var>us-central2</var>"
  node_locations     = ["<var>us-central2-b</var>"]
  cluster            = google_container_cluster.cluster_single_host.name
  initial_node_count = 0
  autoscaling {
    total_min_node_count = 2
    total_max_node_count = 22
    location_policy      = "ANY"
  }

  node_config {
    machine_type = "ct4p-hightpu-4t"
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }
}

替换以下值：

  • your-project：在其中运行工作负载的 Google Cloud 项目。
  • your-node-pool：您要创建的节点池的名称。
  • us-central2：运行工作负载的区域。
  • us-central2-b：运行工作负载的可用区。