为推理工作负载调度 TPU 集合

Trillium (v6e) 包含一项名为“集合调度”的功能,可让您对一组 TPU 切片(单主机或多主机)进行分组,以便提供同一模型的副本。此功能同时适用于 Cloud TPU 和 GKE 配置。

本文档介绍如何将集合调度与 Cloud TPU API 搭配使用。如需详细了解如何将集合调度与 GKE 搭配使用,请参阅 GKE 文档

通过为推理工作负载创建集合,Google Cloud 可以限制并简化对推理工作负载操作的干扰。这对于需要高可用性的推理工作负载非常有用。 Google Cloud 可确保集合的高可用性以管理传入的流量。集合中始终有一部分切片可用于处理传入的流量。

集合中的每个 TPU 切片都具有相同的加速器类型和拓扑。

通过 Cloud TPU API 创建集合

使用 Cloud TPU API 请求已排队的资源时,您可以使用 --workload-type=AVAILABILITY-OPTIMIZED 标志来创建集合。此标志会向 Cloud TPU 基础设施表明,它旨在用于以可用性为重点的工作负载。

以下命令使用 Cloud TPU API 预配集合:

gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
   --project=PROJECT_ID \
   --zone=ZONE \
   --accelerator-type=ACCELERATOR_TYPE \
   --runtime-version=RUNTIME_VERSION \
   --node-count=NODE_COUNT \
   --node-prefix=NODE_PREFIX \
   --workload-type=AVAILABILITY-OPTIMIZED

--node-count 标志用于指定您希望在已排队的资源中包含的切片数。这会创建一组 TPU 切片。

可选:--node-prefix 标志用于指定切片名称的前缀。