Pathways 是一套系統，可用於建立稀疏啟動的大規模多工機器學習 (ML) 系統。這項服務可讓您使用數千或數萬個加速器，並根據處理需求，為不同工作動態分配不同數量的運算資源。
Pathways 可讓單一 JAX 用戶端在大型 TPU 配量 (可能橫跨數千個 TPU 晶片) 之間自動調度管理工作負載，簡化大規模機器學習運算。
Google 內部會使用 Pathways 訓練 Gemini 等大型模型。Pathways on Cloud 為 Google Cloud 客戶帶來同樣的好處。
事前準備
-
請根據您打算如何使用本頁範例，選取對應的分頁標籤：
gcloud
安裝 Google Cloud CLI，然後 使用聯合身分登入 gcloud CLI。 登入後，執行下列指令來初始化 Google Cloud CLI：
gcloud initREST
如要在本機開發環境中使用本頁的 REST API 範例，請使用您提供給 gcloud CLI 的憑證。
安裝 Google Cloud CLI，然後 使用聯合身分登入 gcloud CLI。
詳情請參閱 Google Cloud 驗證說明文件中的「使用 REST 進行驗證」一節。
事前準備
您可以在 Google Kubernetes Engine (GKE) 或使用 Slurm 的 Cluster Director 叢集上執行 Pathways。選取平台的分頁標籤。
GKE
請完成下列前置作業：
本文概述如何在 GKE 上使用 Pathways 代管的 TPU，處理批次、即時和互動式工作負載。並假設您已熟悉如何在 GKE 中使用 TPU (包括 GKE 中的單配量和多配量 TPU)，且具備多配量 TPU 的一般使用經驗。
Cluster Director
請完成下列前置作業：
本文概略說明如何在 Cluster Director 叢集上，搭配 TPU 使用 Pathways 執行批次和互動式工作負載。並假設您已熟悉 Cluster Director、Cluster Director 中的 TPU 機器，以及使用 Slurm 執行工作，且具備多配量 TPU 的一般經驗。
如要建立叢集，請參閱「使用 Pathways 建立 Cluster Director 叢集」。
單一控制器和多控制器
管理及協調多部裝置的運算作業，主要有兩種方式：
| 功能 | 單一控制器 (Pathways) | 多控制器 (JAX 預設) |
|---|---|---|
控管 |
單一控制點：單一用戶端程式做為中央控制器。 |
分散式控制：多個程序參與，每個程序都有自己的 Python 直譯器執行個體。 |
查看 |
統一檢視：用戶端會將所有裝置視為單一統一系統。 |
本機檢視畫面：每個 Python 程序只會看到與其連結的裝置。 |
程式設計 |
簡化程式設計：您只需與單一用戶端互動，系統會顯示為單一大型機器，內含多個本機加速器。 |
單一程式，多筆資料 (SPMD)：主要使用 SPMD 範式，要求所有裝置執行相同的程式。 |
工作彈性 |
支援 SPMD 以外的複雜運算模式，包括非對稱管道平行處理和運算稀疏性。 |
資源管理彈性可能較低，尤其是在不同 TPU 切片之間。 |
Pathways 元件
以下各節將概述 Pathways 架構的主要元件。
Pathways 資源管理工具
這個元件是 Pathways 系統的中央控制層，這項服務會管理所有加速器資源，並協調加速器分配給工作的作業。這項服務會監控 worker 的健康狀態，並處理工作排程、暫停和恢復作業。可做為錯誤和系統狀態的單一聯絡點。這個元件只需要 CPU 資源。
Pathways 用戶端
這個元件是臨時架構執行階段 (IFRT) 的實作，可做為 Pathways 系統的進入點。它會從您的程式接收高階作業 (HLO)。Pathways 用戶端會與 Pathways 資源管理工具協調，根據您的程式碼，決定編譯程式的執行位置。它會向指定的 JAX 用戶端呈現系統的統一檢視畫面。這個元件只需要 CPU 資源。
Pathways 工作站
Pathways worker 程序會在加速器機器 (TPU VM) 上執行。它們會從 IFRT Proxy 伺服器接收程式的編譯執行檔，並在 TPU 上執行運算。Pathways worker 會透過 IFRT Proxy 伺服器，將資料傳回程式。這項元件需要加速器資源。
IFRT Proxy 用戶端
IFRT Proxy 用戶端是臨時架構執行階段 (IFRT) API 的開放原始碼實作，可將使用者程式碼與基礎執行階段分離，並提升程式碼的可攜性和透明度。JAX 會使用這項實作，做為預設多控制器執行階段的替代方案。IFRT Proxy 用戶端是程式與 Pathways 元件之間的通訊橋梁。它會向 IFRT Proxy 伺服器傳送要求，並接收結果。這個元件只需要 CPU 資源。
IFRT Proxy 伺服器
IFRT Proxy 伺服器是 gRPC 伺服器，會接收 IFRT Proxy 用戶端的要求，並轉送至 Pathways 用戶端，由後者處理工作分配。這個元件只需要 CPU 資源。
Sidecar 伺服器
Sidecar 伺服器是 gRPC 伺服器，與 Pathways worker 共同位於加速器 VM 上。這項功能會直接在加速器 VM 上執行使用者指定的 Python 程式碼，減少控制器和加速器之間的資料移轉延遲時間。sidecar 伺服器會透過 gRPC 傳輸上的自訂版本通訊協定，與 Pathways worker 互動。
GKE 上的 Pathways 元件
本節將 Pathways 元件對應至 GKE 元件，例如容器和 Pod。
您可以在下列位置找到 Pathways 容器映像檔：
| 容器類型 | 位置 |
|---|---|
IFRT Proxy 伺服器 |
|
Pathways 資源管理員/工作站 |
|
Pathways 資源管理工具
建立 GKE 叢集後，您可以使用下列 containerSpec 部署 Pathways 資源管理工具：
- name: pathways-rm image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-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 類型和拓撲，格式為tpuTPU_TYPE: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/cloud-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/cloud-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 共用同一個通訊埠。
Cluster Director 上的 Pathways 元件
本節說明如何將 Pathways 元件對應至 Cluster Director 叢集上的 Slurm 工作。
Pathways 指令碼會透過 Enroot 執行 Pathways 元件。setup_env.sh 指令碼會使用下列容器映像檔，這些映像檔與其建構的 Python 環境 JAX 版本 (0.11.1) 相符：
| 容器類型 | 位置 |
|---|---|
IFRT Proxy 伺服器 |
|
Pathways 資源管理員/工作站 |
|
如要啟動 Pathways 元件，請提交 start_pathways.sh 指令碼做為 Pathways 服務工作。您不必自行編寫下列指令：指令碼會從您提交 start_pathways.sh 時設定的環境變數衍生引數，例如 SLICE_SHAPE、MIN_SLICES、MAX_SLICES 和 GCS_BUCKET。
Pathways 資源管理工具
Pathways 服務工作會在一個 CPU 節點上執行 Pathways 資源管理工具。在 Pathways 伺服器容器中執行下列指令：
/pathways/server \ --server_port=29001 \ --node_type=resource_manager \ --instance_type=tpuTPU_TYPE:TOPOLOGY \ --instance_count=MAX_SLICES \ --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \ --enable_gcs_olm=false
這個指令會使用下列引數：
--server_port：Pathways 資源管理工具監聽的通訊埠。IFRT Proxy 伺服器和 Pathways worker 會連線至這個通訊埠。--node_type：節點類型，resource_manager。只有 Pathways 資源管理員會設定這個引數。--instance_type：每個切片的 TPU 類型和拓撲，來自SLICE_SHAPE。例如，SLICE_SHAPE=2x4會設定tpuv6e:2x4。--instance_count：TPU 切片數量，來自MAX_SLICES。--gcs_scratch_location：Cloud Storage bucket 中的路徑，來自GCS_BUCKET，用於編譯快取等暫存檔案。--enable_gcs_olm：Pathways 是否建立 Cloud Storage 物件生命週期規則，清除舊的編譯快取項目。指令碼會關閉這項功能，因為這需要 bucket 的額外權限。
IFRT Proxy 伺服器
Pathways 服務工作會在與 Pathways 資源管理工具相同的 CPU 節點上，執行 IFRT Proxy 伺服器。在 IFRT Proxy 伺服器容器中執行下列指令：
/pathways/proxy_server \ --server_port=29000 \ --resource_manager_address=RM_IP:29001 \ --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \ --num_elastic_instances=MAX_SLICES
這個指令會使用下列引數：
--server_port：IFRT Proxy 伺服器監聽的通訊埠。JAX 程式會連線至這個通訊埠。--resource_manager_address：Pathways 資源管理工具的位址。RM_IP是 CPU 節點的內部 IP 位址。--gcs_scratch_location：與 Pathways 資源管理工具相同的 Cloud Storage 路徑。--num_elastic_instances：在不停止 IFRT Proxy 工作階段的情況下，可無法使用的配量數量，來自MAX_SLICES。工作負載會強制執行MIN_SLICES。
當至少有 MIN_SLICES 個 slice 執行時，服務工作會將 IFRT Proxy 伺服器位址寫入 ~/pathways/cluster.env。JAX 程式 (IFRT Proxy 用戶端) 會讀取這個檔案，連線至 Pathways。JAX 程式會以工作負載 (例如 run_maxtext.sh) 或在 Jupyter 筆記本中執行。
Pathways 工作站
Pathways 服務工作會提交 Slurm 工作陣列 (run_worker_slice.sh)，每個 TPU 切片都有一個陣列任務。每個陣列工作都會在切片中每個 TPU VM 的 Pathways 伺服器容器中，執行下列指令：
/pathways/server \ --server_port=29001 \ --resource_manager_address=RM_IP:29001 \ --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \ --enable_gcs_olm=false
這個指令會使用下列引數：
--server_port：Pathways worker 監聽的通訊埠。Pathways 資源管理工具和 IFRT Proxy 伺服器會連線至這個通訊埠。--resource_manager_address：Pathways 資源管理器的位址，與 IFRT Proxy 伺服器相同。--gcs_scratch_location、--enable_gcs_olm：與路徑資源管理工具相同。
指令碼也會根據切片的節點和陣列工作索引，為每個 worker 設定 TPU 拓撲和 MEGASCALE_* 環境變數。
如果 TPU 切片遭到先占或失敗，Slurm 會重新排入陣列工作。如果叢集節點上有無密碼 sudo，服務工作也會關閉切片節點，以便 Slurm 重新建立節點。Pathways 資源管理工具和 IFRT Proxy 伺服器會持續運作，因此 JAX 程式不需要連線至新位址。