Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

Pathways on Cloud 简介

Pathways 是一种用于构建大规模、多任务且稀疏激活的机器学习 (ML) 系统的架构。它让您能够使用数千或数万个加速器，并能够根据不同任务的处理要求，为它们动态分配不同数量的计算资源。

Pathways 通过使单个 JAX 客户端能够跨大型 TPU 切片（可能跨越数千个 TPU 芯片）编排工作负载，简化了大规模机器学习计算。

Google 在内部使用 Pathways 来训练 Gemini 等大型模型。 Pathways on Cloud 为 Google Cloud 客户带来同样的优势

准备工作

选择标签页以了解您打算如何使用本页面上的示例：
gcloud

安装 Google Cloud CLI，然后使用联合身份登录 gcloud CLI。登录后，运行以下命令来初始化 Google Cloud CLI：
```
gcloud init
```
REST

如需在本地开发环境中使用本页面上的 REST API 示例，请使用您提供给 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 集群上将 Pathways 与 TPU 搭配使用，以处理批量和交互式工作负载。本文档假定您已熟悉 Cluster Director、Cluster Director 中的 TPU 机器以及使用 Slurm 运行作业，并且对多切片 TPU 有一定的了解。

如需创建集群，请参阅使用 Pathways 创建 Cluster Director 集群。

单控制器和多控制器

跨多个设备管理和编排计算主要有两种方式：

功能	单控制器 (Pathways)	多控制器（JAX 默认）
控制	单点控制：单个客户端程序充当中央控制器。	分布式控制：多个进程参与，每个进程都有自己的 Python 解释器实例。
查看	统一视图：客户端将所有设备视为一个统一的系统。	本地化视图：每个 Python 进程只能看到与其连接的设备。
编程	简化编程：您只需与单个客户端交互，系统就像一台拥有许多本地加速器的大型机器。	单程序、多数据 (SPMD)：主要使用 SPMD 范式，要求所有设备运行相同的程序。
灵活性	支持 SPMD 之外的复杂计算模式，包括非对称流水线并行和计算稀疏性。	在资源管理方面可能不太灵活，尤其是在不同的 TPU 切片之间。

Pathways 组件

以下部分概述了 Pathways 架构的主要组件。

Pathways 资源管理器

此组件是 Pathways 系统的中央控制平面。它管理所有加速器资源，并协调作业的加速器分配。它会监控工作器的运行状况，并处理作业调度、暂停和恢复。它充当错误和系统状态的单一联系点。此组件只需要 CPU 资源。

Pathways 客户端

此组件是临时框架运行时 (IFRT) 的实现，用作 Pathways 系统的入口点。它会从您的程序接收高级操作 (HLO)。Pathways 客户端与 Pathways 资源管理器协调，根据您的代码确定将编译后的程序放置在何处以供执行。它向给定的 JAX 客户端呈现系统的统一视图。此组件只需要 CPU 资源。

Pathways 工作人员

Pathways 工作器进程在加速器机器（TPU 虚拟机）上运行。它们从 IFRT 代理服务器接收程序的已编译可执行文件，并在 TPU 上执行计算。Pathways 工作器通过 IFRT 代理服务器将数据发送回您的程序。此组件需要加速器资源。

IFRT 代理客户端

IFRT 代理客户端是临时框架运行时 (IFRT) API 的开源实现，可将用户代码与底层运行时解耦，并提高代码的可移植性和透明度。JAX 使用此实现作为其默认多控制器运行时的替代方案。IFRT 代理客户端充当程序与 Pathways 组件之间的通信桥梁。它会向 IFRT 代理服务器发送请求，并从中接收结果。此组件只需要 CPU 资源。

IFRT 代理服务器

IFRT 代理服务器是一个 gRPC 服务器，它接收来自 IFRT 代理客户端的请求，并将其转发给 Pathways 客户端，后者负责处理工作分配。此组件只需要 CPU 资源。

Sidecar 服务器

边车服务器是一个 gRPC 服务器，与 Pathways 工作器一起位于加速器虚拟机上。它直接在加速器虚拟机上运行用户指定的 Python 代码，以减少控制器和加速器之间的数据传输延迟。边车服务器通过 gRPC 传输层上的自定义版本化协议与 Pathways 工作器进行交互。

架构图，展示了 JAX 客户端、IFRT 代理、Pathways 资源管理器和 Pathways 工作器之间的交互。 — Pathways 组件

GKE 上的 Pathways 组件

本部分将 Pathways 组件映射到 GKE 组件（例如容器和 Pod）。

您可以在以下位置找到 Pathways 容器映像：

容器类型	位置
IFRT 代理服务器	`us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:jax-JAX_VERSION`
Pathways 资源管理器/工作器	`us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:jax-JAX_VERSION`

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 存储桶。

IFRT 代理服务器

您可以使用以下 containerSpec 部署 IFRT 代理服务器：

  - 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：代理服务器用于与 Pathways 资源管理器通信的主机名和端口。该端口必须与用于 Pathways 资源管理器容器的 --server_port 值匹配。
--server_port：IFRT 代理服务器用于与 IFRT 代理客户端通信的端口。
--gcs_scratch_location：用于临时文件的 Cloud Storage 存储桶。

Pathways 工作人员

您可以使用以下 containerSpec 部署 Pathways 工作器：

- 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 Resource Manager 通信的主机名和端口。该端口必须与用于 Pathways 资源管理器容器的 --server_port 值匹配。
--server_port：工作器用于与代理服务器和 Pathways 资源管理器通信的端口。
--gcs_scratch_location：用于临时文件的 Cloud Storage 存储桶。

Pathways 资源管理器、IFRT 代理服务器和 Pathways 工作器可以具有不同的端口，但在本示例中，Pathways 资源管理器和 Pathways 工作器共享同一端口。

Cluster Director 上的 Pathways 组件

本部分将 Pathways 组件映射到 Cluster Director 集群上的 Slurm 作业。

Pathways 脚本会使用 Enroot 运行 Pathways 组件。setup_env.sh 脚本使用以下容器映像，这些映像与它构建的 Python 环境的 JAX 版本 (0.11.1) 相匹配：

容器类型	位置
IFRT 代理服务器	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20261103-jax_0.11.1`
Pathways 资源管理器/工作器	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20261103-jax_0.11.1`

您可以通过将 start_pathways.sh 脚本作为 Pathways 服务作业提交来启动 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 代理服务器和 Pathways 工作器会连接到此端口。
--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 存储桶中的路径，来自 GCS_BUCKET，用于存储编译缓存等临时文件。
--enable_gcs_olm：Pathways 是否创建 Cloud Storage 对象生命周期规则来清理旧的编译缓存条目。该脚本会关闭此功能，因为它需要对存储桶拥有额外的权限。

IFRT 代理服务器

Pathways 服务作业在与 Pathways 资源管理器相同的 CPU 节点上运行 IFRT 代理服务器。它会在 IFRT 代理服务器容器中运行以下命令：

/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 代理服务器监听的端口。您的 JAX 程序会连接到此端口。
--resource_manager_address：Pathways 资源管理器的地址。RM_IP 是 CPU 节点的内部 IP 地址。
--gcs_scratch_location：与 Pathways 资源管理器相同的 Cloud Storage 路径。
--num_elastic_instances：在不停止 IFRT 代理会话的情况下，可以不可用的切片数量，从 MAX_SLICES 开始。您的工作负载强制执行 MIN_SLICES。

当至少 MIN_SLICES 个切片运行时，服务作业会将 IFRT 代理服务器地址写入 ~/pathways/cluster.env。您的 JAX 程序（IFRT 代理客户端）会读取此文件以连接到 Pathways。JAX 程序作为工作负载（例如 run_maxtext.sh）或在 Jupyter 笔记本中运行。

Pathways 工作人员

Pathways 服务作业会提交一个 Slurm 作业数组 (run_worker_slice.sh)，其中每个 TPU 切片对应一个数组任务。每个数组任务都会在其切片中的每个 TPU 虚拟机上的 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 工作器监听的端口。Pathways 资源管理器和 IFRT 代理服务器会连接到此端口。
--resource_manager_address：Pathways 资源管理器的地址，与 IFRT 代理服务器的地址相同。
--gcs_scratch_location、--enable_gcs_olm：与 Pathways 资源管理器相同。

该脚本还会根据切片的节点和数组任务索引，为每个工作器设置 TPU 拓扑和 MEGASCALE_* 环境变量。

如果 TPU 切片被抢占或失败，Slurm 会将其数组任务重新排队。如果您在集群节点上没有密码sudo，则服务作业还会关闭切片的节点，以便 Slurm 可以重新创建这些节点。Pathways 资源管理器和 IFRT 代理服务器会继续运行，因此您的 JAX 程序无需连接到新地址。

Pathways on Cloud 简介 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

准备工作

gcloud

REST

准备工作

GKE

Cluster Director

单控制器和多控制器

Pathways 组件

Pathways 资源管理器

Pathways 客户端

Pathways 工作人员

IFRT 代理客户端

IFRT 代理服务器

Sidecar 服务器

GKE 上的 Pathways 组件

Pathways 资源管理器

IFRT 代理服务器

Pathways 工作人员

Cluster Director 上的 Pathways 组件

Pathways 资源管理器

IFRT 代理服务器

Pathways 工作人员

后续步骤

Pathways on Cloud 简介