GKE 中的 TPU 简介
Google Kubernetes Engine (GKE) 客户现在可以 包含 TPU v4 和 v5e 切片的 Kubernetes 节点池。更多信息 有关 TPU，请参阅系统架构。
使用 GKE 时，您首先需要创建一个 GKE 集群。 然后将节点池添加到 集群。GKE 节点池是 相同的属性。对于 TPU 工作负载，节点池由 TPU 虚拟机组成。
节点池类型
GKE 支持两种类型的 TPU 节点池：
多主机 TPU 切片节点池
多主机 TPU 切片节点池是包含两个或更多个节点池
互连的 TPU 虚拟机。每个虚拟机都连接了一个 TPU 设备。Google Cloud
多主机切片通过高速互连 (ICI) 进行连接。在
多主机切片节点池，您无法向其中添加节点。例如：
您不能先创建 v4-32 节点池，然后再添加额外的 Kubernetes
添加到节点池中。要将一个额外的 TPU 切片添加到
GKE 集群，您必须创建一个新的节点池。
多主机 TPU 切片节点池中的主机被视为一个原子单元。 如果 GKE 无法在切片中部署一个节点， 切片中的节点都会部署
如果需要修复多主机 TPU 切片中的节点，GKE 会 会关停切片中的所有 TPU 虚拟机，强制存储分区中的所有 Kubernetes Pod 删除某些工作负载切片中的所有 TPU 虚拟机启动并运行后， Kubernetes Pod 可以在新切片中的 TPU 虚拟机上调度。
下图显示了一个 v5litepod-16 (v5e) 多主机 TPU 的示例 。此切片有四个 TPU 虚拟机。每个 TPU 虚拟机都连接了四个 TPU v5e 芯片 具有高速互连功能 (ICI)，并且每个 TPU v5e 芯片都有一个 TensorCore。
下图显示了一个 GKE 集群
TPU v5litepod-16 (v5e) 切片（拓扑：4x4）和一个 TPU v5litepod-8 (v5e)
切片（拓扑：2x4）：
如需查看在多主机 TPU 切片上运行工作负载的示例，请参阅 在 TPU 上运行工作负载。
单主机 TPU 切片节点池
单主机切片节点池是指包含一个或多个 独立的 TPU 虚拟机。每个虚拟机都连接了一个 TPU 设备。虽然虚拟机 可以通过数据中心进行通信 网络 (DCN)，则连接到虚拟机的 TPU 未互连。
下图显示了一个具有七个的单主机 TPU 切片的示例
v4-8 机器：
如需查看在单主机 TPU 切片上运行工作负载的示例，请参阅 在 TPU 上运行工作负载。
GKE 节点池的 TPU 机器类型
在创建节点池之前，您需要选择节点池的 TPU 版本和大小，
工作负载所需的 TPU 切片。GKE 支持 TPU v4
标准版 1.26.1-gke.1500 及更高版本、GKE 中的 v5e
标准版 1.27.2-gke.2100 及更高版本，以及 v5p
GKE Standard 版本 1.28.3-gke.1024000 及更高版本。
GKE Autopilot 支持 TPU v4、v5e 和 v5p
1.29.2-gke.1521000 及更高版本。
如需详细了解不同 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 切片机器类型。您的
GKE Standard 集群必须运行控制平面
1.26.1-gke.1500 或更高版本。GKE Autopilot
集群必须运行控制平面 1.29.2-gke.1521000 版或更高版本。
创建 TPU v5p 切片时，使用ct5p-hightpu-4t机器类型
一个主机并包含 4 个芯片。TPU v5p 切片机器类型在
us-west4-a和us-east5-a。GKE Standard
集群必须运行控制平面 1.28.3-gke.1024000 版或更高版本。
GKE Autopilot 必须运行 1.29.2-gke.1521000 或
。如需详细了解 v5p，请参阅 v5p 培训简介。
已知问题和限制
- Kubernetes Pod 数量上限：最多可以运行 256 个 使用单个 TPU 虚拟机中的 Kubernetes Pod。
- 仅限特定预留：在 GKE 中使用 TPU 时，
SPECIFIC是--reservation-affinity标志唯一支持的值gcloud container node-pools create命令的一部分。 - 仅支持抢占式 TPU 的 Spot 虚拟机变体： Spot 虚拟机 与抢占式虚拟机类似， 但无时长上限（24 小时）。
- 不支持费用分配：GKE 费用分配 以及用量计量 不包含任何有关 TPU 使用情况或费用的数据。
- 自动扩缩器可能会计算容量：集群自动扩缩器可能会计算容量 包含 TPU 虚拟机的新节点无法正确计算容量， 可用。然后，集群自动扩缩器可能会执行额外的纵向扩容，并以 创建的节点数超出所需。集群自动扩缩器将缩容 额外节点（如果不需要）。
- 自动伸缩器取消纵向扩容：集群自动伸缩器取消 TPU 的纵向扩容 处于等待状态超过 10 小时的节点池。集群 自动扩缩器稍后会重试此类纵向扩容操作。此行为 降低不使用预留的客户的 TPU 可获取性。
- 污点可能会阻止 TPU 缩容：容忍 如果重新创建节点池，TPU 污点可能会阻止其纵向缩容 排空 TPU 节点池期间的数据。
确保有足够的 TPU 和 GKE 配额
您可能需要增加 创建资源的区域
以下配额具有默认值，可能需要增加：
- Persistent Disk SSD (GB) 配额：每个 Kubernetes 节点的启动磁盘 默认为 100GB因此，此配额应至少设置为 （您预期的最大 GKE 节点数） * 100GB。
- 已用 IP 地址配额：每个 Kubernetes 节点消耗一个 IP 地址。 因此，此配额至少应设为 您预期创建的 GKE 节点。
如需申请增加配额，请参阅申请更高配额。 如需详细了解 TPU 配额的类型，请参阅 TPU 配额。
增加配额的申请可能需要几天时间才能获得批准。如果您 遇到任何难题，在申请增加配额的 请与您的 Google 客户支持团队联系。
迁移 TPU 预留
如果您不打算将现有 TPU 预留与 GKE 中的 TPU 搭配使用， 请跳过本部分并转到创建 Google Kubernetes Engine 集群。
要将预留 TPU 与 GKE 搭配使用，您必须先 将 TPU 预留迁移到基于 Compute Engine 的新预留 系统。
关于此次迁移，您需要注意以下几点重要事项：
- TPU 容量迁移到基于 Compute Engine 的新预留 系统不能与 Cloud TPU Queued Resource API 搭配使用。 如果您打算在预留中使用 TPU 排队资源， 您需要将一部分 TPU 预留迁移到新的 基于 Compute Engine 的预留系统。
- 迁移到 TPU 后，任何工作负载都无法 基于 Compute Engine 的全新预留系统
- 选择执行迁移的时间，并使用 Google Cloud 客户支持团队来安排迁移时间迁移时间范围需要 在工作时间（太平洋时间周一至周五上午 9 点至下午 5 点）。
创建一个 Google Kubernetes Engine 集群
请参阅创建集群 Google Kubernetes Engine 文档。
创建 TPU 节点池
请参阅创建节点池 。
在没有特权模式的情况下运行
如果要缩小容器的权限范围，请参阅 TPU 特权模式。
在 TPU 节点池中运行工作负载
请参阅在 TPU 上运行 GKE 工作负载 。
节点选择器
为了让 Kubernetes 在包含 TPU 虚拟机的节点上调度工作负载， 您必须为 Google Kubernetes Engine 清单中的每个节点指定两个选择器：
- 将
cloud.google.com/gke-accelerator-type设置为tpu-v5-lite-podslice或tpu-v4-podslice. - 将
cloud.google.com/gke-tpu-topology设置为节点的 TPU 拓扑。
训练工作负载和推理工作负载 部分包含说明如何使用这些节点选择器的示例清单。
工作负载安排注意事项
TPU 具有独特的特征，需要特殊的工作负载调度和 Kubernetes 中的管理功能。如需了解详情，请参阅工作负载调度注意事项 。
节点修复
如果多主机 TPU 切片节点池中的节点健康状况不佳， 重新创建整个节点池如需了解详情，请参阅节点自动修复 。
多切片 - 不局限于单个切片
您可以将较小的切片聚合到多切片中以处理较大的切片 训练工作负载。如需了解详情，请参阅 Cloud TPU 多切片。
训练工作负载教程
这些教程重点介绍了如何在多主机 TPU 切片（ 例如 4 台 v5e 机器）。它们涵盖以下模型：
- Hugging Face FLAX 模型：Pokémon 上的 Train Diffusion
- PyTorch/XLA：WikiText 上的 GPT2
下载教程资源
下载每个预训练模型的 Python 脚本和 YAML 规范教程 使用以下命令：
git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git
创建与连接到集群
创建区域级 GKE 集群
控制平面复制到三个可用区，可用性更高
在 us-west4、us-east1 或 us-central2 中创建集群，具体取决于
您正在使用的 TPU 版本。如需详细了解 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
Hugging Face FLAX 模型：Pokémon 上的 Train Diffusion
此示例使用 宝可梦 数据集。
Stable Diffusion 模型是一种潜在的文本到图像模型， 生成逼真的图像。如需详细了解稳定版 扩散，请参阅：
创建 Docker 映像
Dockerfile 位于文件夹下
ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/.
在运行以下命令之前，请确保您的账号拥有正确的 权限 将 Docker 推送到代码库
构建并推送 Docker 映像：
cd ai-on-gke/tutorials-and-examples/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/${project-id}/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
本教程介绍如何使用 HuggingFace 在 v5e TPU 上运行 GPT2 使用 wikitext 数据集构建 PyTorch/XLA 数据集。
创建 Docker 映像
Dockerfile 位于文件夹 ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/ 下。
在运行以下命令之前，请确保您的账号拥有正确的 权限 将 Docker 推送到代码库
构建并推送 Docker 映像：
cd ai-on-gke/tutorials-and-examples/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 的文件中。填写
将图片字段替换为您刚刚创建的图片。
apiVersion: v1
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
volumes:
# Increase size of tmpfs /dev/shm to avoid OOM.
- name: shm
emptyDir:
medium: Memory
# consider adding `sizeLimit: XGi` depending on needs
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/$(project-id)/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'
- 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
volumeMounts:
- mountPath: /dev/shm
name: shm
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
教程：单主机推理工作负载
本教程介绍如何在 Google Cloud 控制台上运行单主机推理工作负载， GKE v5e TPU，用于使用 JAX、TensorFlow、 和 PyTorch。概括来讲，在 GKE 集群：
创建 Cloud Storage 存储桶并设置对该存储桶的访问权限。您使用 Cloud Storage 存储桶用于存储预训练模型。
下载预训练模型，并将其转换为与 TPU 兼容的模型。应用 Kubernetes Pod 下载预训练模型， Cloud TPU Converter 将转换后的模型存储到 Cloud Storage 中， 使用 Cloud Storage FUSE CSI 驱动程序创建存储桶。Cloud TPU Converter 不需要专门的硬件。本教程将向您介绍 模型并在 CPU 节点池中运行 Cloud TPU Converter。
启动转换后模型的服务器。应用 Deployment 该服务使用由 Cloud Storage 中存储的卷支持的服务器框架来提供模型， ReadOnly Many (ROX) 永久性卷Deployment 副本必须运行 在 v5e 切片节点池中，每个节点有一个 Kubernetes Pod。 在 v5e 切片节点池中，每个节点有一个 Kubernetes Pod。
部署负载均衡器以测试模型服务器。该服务器向 LoadBalancer Service 的外部请求。 我们已为 Python 脚本提供了一个示例请求，以测试 模型服务器。
下图展示了负载平衡器如何路由请求。
服务器部署示例
这些示例工作负载采用以下假设进行配置：
- 集群在具有 3 个节点的 TPU v5 节点池中运行
- 节点池使用的是
ct5lp-hightpu-1t机器类型，其中：- 拓扑是 1x1
- TPU 芯片的数量为 1
以下 GKE 清单定义了单个主机 进行部署。
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 字段设置为节点数。
如果您的标准集群运行的是 GKE 1.27 或更低版本， 将以下字段添加到您的清单中：
spec:
securityContext:
privileged: true
您无需在 GKE 中以特权模式运行 Kubernetes Pod 1.28 或更高版本。有关详情，请参阅 在没有特权模式的情况下运行容器。
如果您使用的是其他机器类型：
- 将
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 或更高版本。
在运行requirements.txt
Python 测试脚本。
如果您没有在本地环境中正确设置 Python，则可以 使用 Cloud Shell 下载并运行 Python 脚本。
设置集群
使用
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
示例工作负载假定：
- 您的集群正在使用具有 3 个节点的 TPU v5e 节点池运行。
- TPU 节点池使用的是机器类型“
ct5lp-hightpu-1t”。
如果您使用的集群配置与上述集群配置不同， 您需要修改服务器部署清单。
对于 JAX Stable Diffusion 演示，您需要一个节点池
具有 16 Gi+ 可用内存的机器类型（例如 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 存储分区的访问权限。
为您的应用创建 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 服务账号。IAM 服务账号 必须与 Cloud Storage 存储桶位于同一项目中。
确保您的 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 服务账号。
允许 Kubernetes 服务账号模拟 IAM 为服务账号添加 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 服务账号。your-project-id：您在其中创建了 GKE 集群。您的 Cloud Storage 存储分区和 GKE 集群可以位于同一项目中，也可以位于不同的项目中。
使用 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 服务账号。
运行以下命令，在 YAML 文件中填充存储桶的名称 本演示：
find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"
将
your-bucket-name替换为 Cloud Storage 存储分区的名称。使用以下命令创建永久性卷和永久性卷声明 命令：
kubectl apply -f pvc-pv.yaml
JAX 模型推断和传送
安装 Python 依赖项，以便运行 发送到 JAX 模型服务的请求。
pip install -r jax/requirements.txt
运行 JAX BERT E2E 服务演示：
此演示使用预训练的 BERT 模型 来自 Hugging Face。
Kubernetes Pod 会执行以下步骤：
- 下载并使用示例中的 Python 脚本
export_bert_model.py用于将预训练的 BERT 模型下载到临时目录的资源。 - 使用 Cloud TPU Converter 图片将预训练模型从 并将模型存储在您所需的 Cloud Storage 存储桶中， 在设置过程中创建的 ID。
此 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 可能需要一分钟时间来读取 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 超过 16Gi 的可用内存，如 设置集群。
Kubernetes Pod 会执行以下步骤：
- 从以下位置下载并使用 Python 脚本
export_stable_diffusion_model.py： 可下载预训练稳定扩散模型的示例资源 临时目录 - 使用 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 可能需要一分钟时间来读取 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 模型服务。
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 可能需要一分钟时间来读取 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
PyTorch 模型推断和服务
安装 Python 依赖项，以便运行 发送到 PyTorch 模型服务：
pip install -r pt/densenet161/requirements.txt
运行 TorchServe Dense161 E2E 服务演示：
生成模型归档。
- 应用模型归档：
kubectl apply -f pt/densenet161/model-archive.yml
- 使用以下命令验证模型是否已正确安装：
kubectl get pods densenet161-model-archive
STATUS可能需要几分钟时间才能读取Completed。使用 TorchServe 提供模型：
应用模型服务部署：
kubectl apply -f pt/densenet161/deployment.yml
使用以下命令验证服务器是否按预期运行：
kubectl get deployment densenet161-deployment
AVAILABLE可能需要一分钟时间来读取3。应用负载均衡器服务：
kubectl apply -f pt/densenet161/loadbalancer.yml
使用 以下命令：
kubectl get svc densenet161-service
EXTERNAL_IP可能需要几分钟的时间才能列出 IP。
向模型服务器发送测试请求：
从负载均衡器获取外部 IP：
EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')运行测试请求脚本，以将请求 (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
通过运行以下
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 问题。
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 上调度它们
具有指定 TPU 拓扑（例如 v4-32）的切片。另外，假设您
GKE 集群中有两个 v4-32 TPU 切片；我们将
将这两个切片称为切片 X 和 Y。由于您的集群有足够的容量
同时安排两个作业，理论上两个作业应该都会很快被安排 - 其中一个作业同时
两个 TPU v4-32 切片中的每一个。
不过，如果没有仔细规划，就有可能制定一个时间表， 死锁。假设 Kubernetes 调度器从 Job 调度一个 Kubernetes Pod A 在切片 X 上，然后在切片 X 上安排作业 B 中的一个 Kubernetes Pod。在 在这种情况下，如果作业 A 的 Kubernetes Pod 亲和性规则，调度器将 尝试在切片 X 上为作业 A 调度所有剩余的 Kubernetes Pod。对 作业 B。因此作业 A 和作业 B 都无法在 单个切片。这会导致调度死锁。
为了避免调度死锁的风险，您可以使用 Kubernetes 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 = "${project-id}"
name = "${node-pool-name}"
location = "${location}"
node_locations = ["${node-locations}"]
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 = ["${reservation-name}"]
}
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 = "${project-id}.svc.id.goog"
}
addons_config {
gcs_fuse_csi_driver_config {
enabled = true
}
}
}
resource "google_container_node_pool" "single_host_tpu" {
provider = google-beta
project = "${project-id}"
name = "${node-pool-name}"
location = "${location}"
node_locations = ["${node-locations}"]
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：运行工作负载的可用区。