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 系統的中央控制層，這項服務會管理所有加速器資源，並負責協調使用者工作的加速器分配作業。負責監控工作人員的健康狀態，並處理工作排程、暫停和繼續作業。可做為錯誤和系統狀態的單一聯絡點。這項元件只需要 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 工作人員互動。

PathwaysJob API

PathwaysJob API 是 OSS Kubernetes 原生 API，可用於部署 ML 訓練和批次推論工作負載。PathwaysJob 的控制器會運用 JobSet API 管理所有 Pathways 元件的生命週期和協調作業。這個自訂資源定義 (CRD) 提供高階介面，可定義 Pathways 工作負載，並在常見情境中抽象化，不必直接管理個別 Pod 規格。如需所有參數的完整清單及其具體含義，請參閱 GitHub 上的 PathwaysJob API 說明文件。

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

下表說明 PathwaysJob API 的設定：

屬性	說明
`apiVersion`	指定 `PathwaysJob` 自訂資源定義 (CRD) 的 API 版本：pathways-job.pathways.domain/v1。
`kind`	將 Kubernetes 物件識別為 `PathwaysJob`。
`metadata.name`	Kubernetes 中的 `PathwaysJob` 物件名稱，通常遵循 pathways-. 模式。
`spec`	定義 `PathwaysJob` 的所需狀態和設定。
`spec.maxRestarts`	如果 `PathwaysJob` 發生故障，系統自動重新啟動的次數上限。
`spec.pathwaysVersion`	(選用) 指定要在這項工作的 Pathways 環境中使用的 JAX 架構版本 (例如 jax-0.5.3)。
`spec.workers`	定義 `PathwaysJob` 工作站集區設定的陣列，通常會使用 TPU 資源。
`spec.workers[].type`	要用於工作站節點的 TPU 機器類型 (例如，$TPU_MACHINE_TYPE 可以是 ct6e-standard-4t)
`spec.workers[].topology`	分配給工作人員的 TPU 配量拓撲 (例如，$TOPOLOGY 可以是 2x2、4x4、2x2x2)。
`spec.workers[].numSlices`	要為工作站集區佈建的 TPU 節點數量 (例如，$WORKLOAD_NODEPOOL_COUNT 可以是 2)。
`spec.workers[].maxSliceRestarts`	(選用) 如果切片中的個別工作站失敗，可以重新啟動的次數上限。
`spec.customComponents`	(選用) 陣列，可讓您定義及部署自訂元件 (例如 Proxy 伺服器、Pathways 伺服器或其他工作站)，與主要工作並行。
`spec.customComponents[].componentType`	指定要定義的自訂元件類型 (例如 proxy_server、pathways_server、worker、colocated_python_sidecar)。
`spec.customComponents[].image`	這個自訂元件的容器所要使用的 Docker 映像檔。
`spec.customComponents[].customFlags`	啟動容器時要傳遞至容器的自訂指令列標記陣列。
`spec.customComponents[].customEnv`	要在容器內設定的自訂環境變數陣列。每個元素都有名稱和值。
`spec.pathwaysDir`	`PathwaysJob` 用於儲存編譯構件和其他暫時資料的 Cloud Storage 值區。執行工作負載前，您必須先建立這個值區。
`spec.controller`	路徑控制器設定，用於管理整體工作執行作業。
`spec.controller.deploymentMode`	指定 Pathways 控制器的 CPU 資源 (Pathways 資源管理工具和 Proxy 伺服器) 部署方式。預設模式會將這些節點部署到專屬 CPU 節點，而 colocate_head_with_workers 則會將節點部署到 TPU 工作站。
`spec.controller.elasticSlices`	(選用) 工作執行期間可停機的 TPU 節點數量上限，超過此上限即視為不正常。
`spec.controller.template`	(選用) 定義使用者作業的 Pod 範本。批次工作負載必須提供這項資訊，但互動式工作負載則不必。
`spec.controller.template.spec`	使用者工作的 Pod 規格。
`spec.controller.template.spec.containers`	定義要在使用者工作中執行的容器陣列
`spec.controller.template.spec.containers[].name`	使用者作業中的容器名稱 (在本範例中為 main)。
`spec.controller.template.spec.containers[].image`	主要容器中要用於容器的 Docker 映像檔 (在本範例中為 python:3.11)。
`spec.controller.template.spec.containers[].command`	主要容器啟動時要執行的指令。在本範例中，指令碼會安裝 `pathwaysutils`、初始化 Pathways，並列印 JAX 裝置。

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>`

路徑資源管理員

建立 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：用於暫時檔案的 GCS 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 伺服器和路徑工作站可以有不同的通訊埠，但在這個範例中，路徑資源管理工具和路徑工作站共用同一個通訊埠。

Cloud 上的路徑簡介 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

事前準備

單一控制器和多個控制器

路徑元件

路徑資源管理員

路徑用戶端

Pathways 工作者

IFRT Proxy 用戶端

IFRT Proxy 伺服器

Sidecar 伺服器

PathwaysJob API

GKE 上的 Pathways 元件

路徑資源管理員

IFRT Proxy 伺服器

Pathways 工作者

後續步驟

Cloud 上的路徑簡介