Pathways on Cloud 简介

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

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

Pathways 在 Google 内部用于训练 Gemini 等大型模型。Cloud 上的 Pathway 可为 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 资源。

Pathways 客户端

这是临时框架运行时 (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 进行交互。

显示 Pathways 组件之间的关系。
开发者在线课程组件

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：用于临时文件的 Cloud Storage 存储桶。

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 工作器共享同一端口。

后续步骤