Pathways on Cloud 简介

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

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

Pathways 在 Google 内部用于训练 Gemini 等大型模型。“云端之路”计划同样能为 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 系统的中央控制平面。它管理所有加速器资源，并负责协调用户作业的加速器分配。它会监控工作器的健康状况，并处理作业的调度、暂停和恢复。它可作为错误和系统状态的单一联系点。此组件仅需要 CPU 资源。

Pathway 客户端

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

Pathways worker

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

IFRT 代理客户端

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

IFRT 代理服务器

此 gRPC 服务器从 IFRT 代理客户端接收请求，并将其转发给 Pathways 客户端，后者负责实际的工作分配。此组件仅需要 CPU 资源。

Sidecar 服务器

此 gRPC 服务器与加速器虚拟机上的 Pathways 工作器位于同一位置，可直接在加速器虚拟机上运行用户指定的 Python 代码，从而缩短从控制器到加速器的数据传输延迟时间。Sidecar 服务器通过 gRPC 传输上的自定义版本化协议与 Pathways worker 进行交互。

PathwaysJob API

PathwaysJob API 是一种 OSS Kubernetes 原生 API，可用于部署机器学习训练和批量推理工作负载。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`	（可选）一个数组，用于定义和部署自定义组件（例如代理服务器、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`	用于管理整个作业执行过程的 Pathways 控制器配置设置。
`spec.controller.deploymentMode`	指定如何部署 Pathways 控制器的 CPU 资源（Pathways 资源管理器和代理服务器）。默认模式会将它们部署在专用 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 组件

本部分将 Pathway 组件映射到 Google Kubernetes Engine 组件，例如容器和 pod。

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

容器类型	位置
IFRT 代理服务器	`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 存储桶。

IFRT 代理服务器

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

 - 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：代理服务器用于与 Pathways 资源管理器通信的主机名和端口。该端口应与用于 Pathways 资源管理器容器的 --server_port 值相同。
--server_port：IFRT 代理服务器使用此端口与 IFRT 代理客户端通信。
--gcs_scratch_location：用于临时文件的 GCS 存储桶。

Pathways worker

您可以使用以下 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：工作器使用此端口与代理服务器和 Pathways 资源管理器进行通信。
--gcs_scratch_location：用于临时文件的 Cloud Storage 存储桶。

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

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

准备工作

单控制器和多控制器

Pathways 组件

Pathways 资源管理器

Pathway 客户端

Pathways worker

IFRT 代理客户端

IFRT 代理服务器

Sidecar 服务器

PathwaysJob API

GKE 上的 Pathways 组件

Pathways 资源管理器

IFRT 代理服务器

Pathways worker

后续步骤

Pathways on Cloud 简介