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 的一般使用經驗

單一控制器和多控制器

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

功能

單一控制器 (Pathways)

多控制器 (JAX 預設)

控管

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

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

查看

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

本機檢視：每個 Python 程序只會看到與其連結的裝置。

程式設計

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

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

工作彈性

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

資源管理彈性可能較低，尤其是在不同 TPU 切片之間。

Pathways 元件

以下章節將概述 Pathways 架構的主要元件。

Pathways 資源管理工具

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

Pathways 用戶端

這是臨時架構執行階段 (IFRT) 的實作，可做為 Pathways 系統的進入點。它會從程式接收高階作業 (HLO)。Pathways 用戶端負責與 Pathways 資源管理工具協調，根據使用者程式碼決定編譯程式的執行位置。它會向特定 JAX 用戶端呈現系統的整合式資料檢視。這個元件只需要 CPU 資源。

Pathways 工作人員

這些程序會在加速器機器 (TPU VM) 上執行。它們會從 IFRT Proxy 伺服器接收程式的編譯執行檔，並在 TPU 上執行運算。Pathways worker 會透過 IFRT Proxy 伺服器，將資料傳回程式。這項元件需要加速器資源。

IFRT Proxy 用戶端

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

IFRT Proxy 伺服器

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

Sidecar 伺服器

這個 gRPC 伺服器與 Pathways worker 位於加速器 VM 的相同位置，可直接在加速器 VM 上執行使用者指定的 Python 程式碼，減少從控制器到加速器的資料移轉延遲時間。sidecar 伺服器會透過 gRPC 傳輸上的自訂版本化通訊協定，與 Pathways worker 互動。

顯示 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 部署 Pathways 資源管理工具：

  - 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：Pathways 資源管理工具會使用這個通訊埠，與其他 Pathways 元件通訊。
  • --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 worker：

- 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：worker 會使用這個通訊埠與 Proxy 伺服器和 Pathways 資源管理工具通訊。
  • --gcs_scratch_location：用於暫存檔案的 Cloud Storage bucket。

Pathways 資源管理工具、IFRT Proxy 伺服器和 Pathways worker 都可以使用不同的通訊埠，但在這個範例中，Pathways 資源管理工具和 Pathways worker 共用同一個通訊埠。

後續步驟