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

Trillium (v6e) 包含一项名为“集合调度”的功能,可让您对一组 TPU 切片(单主机或多主机)进行分组,以便提供同一模型的副本。此功能同时适用于 Cloud TPU 和 Google Kubernetes Engine (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 标志用于指定切片名称的前缀。