Run Ray on TPU, Part 1: The foundations

Google Developers Blog

Google Developers Blog2026年7月21日

Run Ray on TPU, Part 1: The foundations

8.5内容质量

TL;DR · AI 摘要

Ray 2.55正式支持TPU加速，通过GKE和Ray Core协作实现TPU资源调度，开发者可复用现有API进行分布式训练。

核心要点

Ray 2.55版本将TPU纳入官方加速器支持体系，提供预构建镜像和全栈库支持
TPU slice必须作为整体单元调度，类似GPU的多卡箱结构，依赖ICI高速互联
GKE负责TPU拓扑标签管理，Ray Core通过标签识别并锁定完整slice资源

结构提纲

按章节快速跳转。

§引言
宣布Ray 2.55正式支持TPU加速，开发者可复用现有API进行分布式训练。
·Ray与TPU基础概念
解释Ray框架如何将TPU视为可调度资源，强调slice结构与GPU的异同。
›TPU slice架构
详细说明TPU slice的硬件组网方式及其对分布式训练的影响。
·GKE与Ray Core协作机制
描述GKE如何管理TPU拓扑标签，Ray Core如何锁定完整slice资源。
›开发实践
展示现有Ray代码无需修改即可运行于TPU的实现路径。

思维导图

用一张图看清主题之间的关系。

查看大纲文本（无障碍 / 无 JS 友好）

Run Ray on TPU
- TPU架构
  - slice结构
  - ICI互联
- Ray集成
  - 任务调度
  - Actor模型
- GKE管理
  - 拓扑标签
  - 资源锁定

金句 / Highlights

值得收藏与分享的关键句。

TPU slice必须作为整体单元调度，类似GPU的多卡箱结构，依赖ICI高速互联
— 第3段
⬇︎ 下载 PNG 𝕏 分享到 X
GKE负责TPU拓扑标签管理，Ray Core通过标签识别并锁定完整slice资源
— 第5段
⬇︎ 下载 PNG 𝕏 分享到 X
Ray 2.55版本提供预构建镜像和全栈库支持，替代原有实验性容器方案
— 第2段
⬇︎ 下载 PNG 𝕏 分享到 X

#Ray#TPU#GKE#分布式计算#AI加速器

打开原文

在 TPU 上运行 Ray，第 1 部分：基础 - Google Developers Blog

Google Tag Manager (noscript)

End Google Tag Manager (noscript)

HTML

在 TPU 上运行 Ray，第 1 部分：基础

2026 年 7 月 20 日

Ivan Nardini

AI 开发者关系

Facebook

Twitter

LinkedIn

邮件

TL;DR：如果你已经使用 Ray 在 GPU 上扩展 Python，现在你的代码也可以在 TPU（Tensor Processing Unit，Google 的 AI 加速芯片）上运行，并且拥有你已经熟悉的官方 API 全面支持。任务和演员模型、一个 JaxTrainer、相同的 Ray Serve 部署，只需指向由 Google Kubernetes Engine（GKE）管理的 TPU。

从 Ray 2.55 版本开始，Google Cloud TPU 成为 Ray 的一级加速器。这意味着 TPU 现在已纳入 Ray 的发布流程，拥有官方预构建镜像和核心库的全面支持，而不是过去需要自行构建容器并依赖社区帮助的“实验性”路径。在本系列“在 TPU 上运行 Ray”中，你将学习 TPU 的一个切片（slice）对 Ray 来说只是另一个可调度的资源（第 1 部分），然后逐步了解每个库（第 2 部分）。

Ray 与 TPU：快速入门

Ray 是一个分布式计算框架：你编写 Python 代码，Ray 会将其在集群中作为任务（无状态函数）和演员（有状态工作节点）运行。对 Ray 来说，TPU 只是另一个可调度的资源，就像 GPU 一样。你提出需求，Ray 会将你的工作分配到 TPU 上。

但有一点需要注意，之后我们再继续。

TPU 芯片通过专用高速链路（称为 ICI，芯片间互连）连接成固定组，称为一个切片（slice）：多个主机机器（虚拟机）的芯片共享这个链路。多主机模型必须部署在完整的切片上，否则工作节点无法互相通信，任务会陷入停滞。

如果你用 GPU 思考，可以想象一个切片就像一个包含多个 GPU 的盒子，其中的高速互连（NVLink）仅存在于盒子内部。如果将工作节点分布在两个没有连接的盒子中，集体操作（如同步梯度的 all-reduce 步骤）将永远无法完成，训练会陷入停滞。TPU 切片的行为方式相同：ICI 就是那条连接线，它仅连接一个切片的芯片。

这就是“在 TPU 上运行 Ray”需要特殊处理的全部原因。必须确保所有工作节点都落在一个完整的切片上。在 GPU 上你几乎不需要考虑这一点；但在 TPU 上，这是关键，而 Ray 和 GKE（Google Kubernetes Engine，Google 的托管 Kubernetes）会为你处理。

你还会在各处看到“拓扑”（topology）这个词：它描述切片的形状，例如 16 芯片切片写成 4x4。你请求的是拓扑，而不是芯片数量。

一旦理解了 TPU 的切片和拓扑概念，现有的 Ray 堆栈和开发流程将保持不变，并可在 GKE 提供的 TPU 切片上运行。下图用一张图展示了整个系统：左侧是你编写的代码（你已经使用的 Ray 库）；中间是 Ray Core 层，它负责预留完整的切片；右侧是 GKE 管理层，它提供硬件并进行标记，使 Ray 能够识别切片边界。

GKE 提供一个切片并标记其主机，Ray Core 读取这些标签以一次性预留整个切片，而你的库调用则位于顶层，声明拓扑并无需其他操作。任何地方都不需要手动编写放置代码。本部分的其余内容将介绍底层的两个层级：GKE 然后是 Ray Core，第 2 部分则涵盖 Ray AI 库。

GKE 如何在 TPU 上调度 Ray

通过 GKE 的 Ray Operator 插件，你可以在 TPU 上运行 Ray。

code

# Autopilot（全托管节点）
gcloud container clusters create-auto CLUSTER \
  --enable-ray-operator --location=LOCATION

# 或 Standard（您管理节点池）
gcloud container clusters create CLUSTER \
  --addons=RayOperator --location=LOCATION

Shell

已复制

该单一标志安装了对TPU至关重要的两个组件。第一个是 KubeRay，这是一个 Kubernetes 操作器，它将 RayCluster、RayService 和 RayJob 的 YAML 转换为运行中的 Ray 集群；它与您在 GPU 上使用的 KubeRay 是相同的。第二个是特定于 TPU 的部分：Ray TPU webhook，它为每个 TPU 主机打上类似 ray.io/tpu-slice-name 的标签，使 Ray 能够识别哪些机器连接到同一片。这个标签是整个系统运作的关键线索。

从那里开始，您通过声明与请求任何节点相同的方式来请求 TPU，使用节点选择器指定生成版本和拓扑结构，并将芯片数量作为资源。多主机片增加了一个字段：numOfHosts。

code

# 在 RayCluster workerGroupSpec 内部
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice   # TPU 生成版本
  cloud.google.com/gke-tpu-topology: "4x4"              # 片形状
# ... 通过 google.com/tpu 资源限制请求芯片
numOfHosts: 4   # 多主机：组成此片的主机虚拟机数量

GKE 会分配该片，webhook 会为其打标签，Ray 读取这些标签。您编写 Python 代码。一旦附加组件就绪，您将看到 KubeRay 操作器 pod 正在运行，应用该清单会启动一个头 pod 和每个片主机的一个工作 pod。get-started 示例的集群步骤通过 Terraform 会为所有这些内容提供支持。

实际上将您的工作节点保持在一起的是位于此层上方的 Ray Core 原语——片放置组，而指南的其余部分正是从此处开始。

TPU 上的 Ray Core

Ray Core 是基础层，所有其他内容都构建在任务和演员引擎及调度器之上。其 TPU 支持位于公共的 ray.util.tpu API 中，真正需要了解的只有一个函数：slice_placement_group()。它将前面提到的“将我的工作节点保留在一个完整的片上”的保证转化为单个调用，通过匹配 webhook 标签来原子性地保留整个片（所有主机或无）。

code

from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

# 原子性地保留一个完整的 v6e 4x4 片（4 台主机上的 16 个芯片）
spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.ready(), timeout=600)

@ray.remote(resources={"TPU": 4})
def worker(rank, world): ...

tasks = [
    worker.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=spg.placement_group)
    ).remote(rank=i, world=spg.num_hosts)
    for i in range(spg.num_hosts)
]

Python

需要强调的是，您很少直接调用 slice_placement_group。Ray AI 库（Data、Train、Serve）会为您调用它，因此在实际操作中您只需声明拓扑结构，它们会处理片分配。您只有在编写自定义分布式工作负载（非 Train、Serve 或 Data）时才需要直接使用 slice_placement_group()。一个需要注意的细节是：该 API 是公开的但标记为 alpha（@PublicAPI(stability="alpha")），因此目前可用，但不同版本之间接口仍可能变化。

这是基础。接下来是库。

你现在对整个概念模型已经有了全面的理解：切片必须保持完整，GKE 会为其分配资源并打标签，而 Ray Core 会将其作为一个整体保留下来，因此你永远不需要手动编写放置代码。你实际构建的所有内容都基于这一机制并复用它。

在第二部分中，我们将探讨如何在 TPU 上使用 Ray AI 库，通过 vLLM 为 LLM 提供服务，利用 Ray Data 为切片提供数据，并使用 JaxTrainer 进行训练。

其他资源

可运行示例：在 Kubernetes Engine 示例中启动 Ray on TPU，包含 cluster、serve、data 和 train 在单个 v6e 切片上的演示。

GKE 上的 Ray Operator 附加组件

在 KubeRay 中使用 TPU

目前，感谢你的阅读！如果你有任何其他问题或反馈，欢迎在社交媒体上联系我（LinkedIn、X）。

愉快地构建！

发布于：

AI

操作指南

公告

AI

云服务

公告

解决方案

在 Gemini 企业代理平台中扩展选择：引入并行网络搜索的 Grounding 功能

2026年7月16日

操作指南

最佳实践

为什么我们构建了 ADK 2.0

2026年7月1日

行业趋势

使用模块化提示转换构建可扩展的 AI 代理

从你的编码代理驱动代理质量飞轮

2026年6月30日

导航点 /