推論工作負載的 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 標記會指定切片名稱的前置字串。