Cloud 上的路徑簡介

Pathways 是一套系統，可用於建立稀疏啟動的大規模多工機器學習系統。可使用數千或數萬個加速器，並根據不同工作的處理需求，動態分配不同數量的運算資源。

Pathways 可讓單一 JAX 用戶端在多個大型 TPU 切片 (可能涵蓋數千個 TPU 晶片) 中調度工作負載，簡化大規模機器學習運算。

Google 內部會使用 Pathways 訓練 Gemini 等大型模型。 Pathways on Cloud 也能為 Google Cloud 客戶帶來同樣的優勢。

事前準備

請確認您已備妥以下項目：

本文提供總覽，說明如何在 Google Kubernetes Engine (GKE) 上使用 Pathways 管理的 TPU，處理批次、即時和互動式工作負載。本文假設您已熟悉如何搭配使用 TPU 和 GKE，包括在 Google Kubernetes Engine 上使用單一和多個 TPU 切片，以及多個 TPU 切片的一般使用體驗。

單一控制器和多個控制器

管理及自動化調度多部裝置的運算作業，主要有兩種不同的方式：

功能

單一控制器 (路徑)

多重控制器 (JAX 預設)

控管

單一控制點：單一用戶端程式會做為中央控制器。

分散式控制：多個程序參與，每個程序都有自己的 Python 解譯器執行個體。

查看

整合式資料檢視：客戶會將所有裝置視為單一整合系統。

本地化檢視畫面：每個 Python 程序只會看到連線至該程序的裝置。

程式設計

簡化程式設計：使用者與單一用戶端互動，讓系統看起來像是一部大型機器，內含許多本機加速器。

SPMD：主要使用 SPMD 範例，要求所有裝置執行相同程式。

工作彈性

支援 SPMD 以外的更複雜運算模式，包括非對稱管線平行處理和運算稀疏性。

資源管理彈性較低，尤其是在不同 TPU 節點之間。

學習路徑元件

以下各節將說明 Pathways 架構的主要元件。

Pathways 資源管理員

這是 Pathways 系統的中央控制層。負責管理所有加速器資源，並協調使用者工作加速器的分配作業。負責監控工作人員的健康狀態，並處理工作排程、暫停和繼續作業。這個單一聯絡點可處理錯誤和系統狀態。這項元件只需要 CPU 資源。

路徑用戶端

這是「過渡架構執行階段」(IFRT) 的實作項目，可做為 Pathways 系統的進入點。並接收來自您計畫的高階作業 (HLO)。Pathways 用戶端負責與 Pathways 資源管理工具協調，根據使用者程式碼決定要將編譯後的程式放在何處執行。並向特定 JAX 用戶端提供系統的整合式資料檢視。這項元件只需要 CPU 資源。

Pathways 工作人員

這些程序會在加速器機器 (TPU VM) 上執行。這些裝置會從 IFRT Proxy 伺服器接收程式的已編譯可執行檔，並在 TPU 上執行運算作業。路徑工作人員會透過 IFRT Proxy 伺服器，將資料傳回您的程式。這個元件需要加速器資源。

IFRT Proxy 用戶端

這是 Interim Framework Runtime (IFRT) API 的 OSS 實作，可將使用者程式碼與基礎執行階段分離，並提升程式碼可攜性和透明度。JAX 會使用這項實作做為預設多控制器執行階段的替代方案。IFRT Proxy 用戶端是程式與 Pathways 元件之間的通訊橋樑。並向 IFRT Proxy 伺服器傳送要求，以及接收結果。這是 IFRT API 的 OSS 實作項目。這項元件只需要 CPU 資源。

IFRT Proxy 伺服器

這個 gRPC 伺服器會接收 IFRT 代理用戶端的要求，並轉送至 Pathways 用戶端，由後者處理實際的工作分配作業。這項元件只需要 CPU 資源。

Sidecar 伺服器

這個 gRPC 伺服器與加速器 VM 上的 Pathways 工作站位於同一位置，可直接在加速器 VM 上執行使用者指定的 Python 程式碼，減少從控制器到加速器的資料傳輸延遲。在 gRPC 傳輸上，側車伺服器會透過自訂版本化通訊協定與 Pathways 工作人員互動。

顯示 Pathways 元件的關係。
路徑元件

GKE 上的 Pathways 元件

本節會將 Pathways 元件對應至 Google Kubernetes Engine 元件，例如容器和 Pod。

您可以在下列位置找到 Pathways 容器映像檔。

容器類型

位置

IFRT Proxy 伺服器

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>

Pathways 資源管理員/工作人員

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

Pathways 資源管理員

建立 GKE 叢集後，您可以使用下列 containerSpec 部署路徑資源管理工具：

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

引數說明：

  • --server_port：路徑資源管理工具會使用這個通訊埠與其他路徑元件通訊。
  • --node_type：節點類型。如果是 Pathways 資源管理工具，這項屬性應設為「resource_manager」，其他容器則不需要。
  • --instance_count：TPU 配量數量。
  • --instance_type：配量的 TPU 類型和拓撲。格式為 tpu{TPU type}:{TPU topology}，例如 tpuv5e:4x4。
  • --gcs_scratch_location：用於暫存檔案的 Cloud Storage bucket。

IFRT Proxy 伺服器

您可以使用下列 containerSpec 部署 IFRT Proxy 伺服器：

 - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

引數說明：

  • --resource_manager_address：Proxy 伺服器用來與 Pathways 資源管理工具通訊的主機名稱和通訊埠。這個連接埠應與 Pathways 資源管理工具容器使用的 --server_port 值相同。
  • --server_port：IFRT Proxy 伺服器會使用這個通訊埠與 IFRT Proxy 用戶端通訊。
  • --gcs_scratch_location：用於暫存檔案的 Cloud Storage bucket。

Pathways 工作人員

您可以使用下列 containerSpec 部署 Pathways 工作人員：

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

引數說明：

  • --resource_manager_address：TPU 工作人員用來與 Pathways 資源管理工具通訊的主機名稱和通訊埠。通訊埠應與 Pathways 資源管理工具容器使用的 --server_port 值相同
  • --server_port：工作人員會使用這個通訊埠與 Proxy 伺服器和 Pathways 資源管理工具通訊。
  • --gcs_scratch_location：用於暫存檔案的 Cloud Storage bucket。

路徑資源管理工具、IFRT Proxy 伺服器和路徑工作站可以有不同的通訊埠，但在這個範例中，路徑資源管理工具和路徑工作站共用同一個通訊埠。

後續步驟