Pathways on Cloud 简介

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

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

Pathways 在 Google 内部用于训练 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 系统的中央控制平面。它管理所有加速器资源，并负责协调用户作业的加速器分配。它会监控工作器的运行状况，并处理作业调度、暂停和恢复。它充当错误和系统状态的单一联系点。此组件只需要 CPU 资源。

Pathways 客户端

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

Pathways 工作器

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

IFRT 代理客户端

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

IFRT 代理服务器

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

Sidecar 服务器

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

显示 Pathways 组件的关系。
Pathways 组件

GKE 上的 Pathways 组件

本部分将 Pathways 组件映射到 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 工作器

您可以使用以下 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 工作器共享同一端口。

后续步骤