Google 详解 Ray AI 库在 TPU 上的运行：Serve、Data 与 Train 抽象化多主机调度与分布式训练

简要总结：第2部分，总共2部分。第1部分：https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ 介绍了你需要的一个硬件概念和下面的两层（GKE 和 Ray Core）。本部分展示了你实际构建时使用的库：Ray Serve、Ray Data 和 Ray Train。

如果你是第一次看到这里，快速回顾一下。将 Ray 运行在 TPU 上归结为一个注意事项：TPU 芯片被固定分组称为切片（主机虚拟机共享一个称为 ICI 的高速链接），多主机模型必须落在一个完整的切片上，否则其工作节点无法互相访问，作业会挂起。

配备 Ray Operator 插件的 Google Kubernetes Engine (GKE) 会提供切片并标记其主机，Ray Core 原语 slice_placement_group() 可以一次性预留整个切片。你声明一个拓扑（切片形状，例如 16 块芯片的 4x4），下面的库会为你处理放置问题。

有了 Core 在底层处理放置，所有库都遵循同样的模式：声明一个拓扑，让 Core 预留切片。各个库不同的只是你声明的对象。我们将按大多数团队采用的顺序介绍，从 serving 开始。

大多数团队从 Serving 开始。一种需要多个 GPU 才能适配的模型可以运行在单个 TPU 主机上，并且 TPU 通常在推理中更可用且性价比高。Ray Serve 为你提供常见的自动伸缩、负载均衡和多模型组合，并在 TPU 上通过 vLLM（一个高吞吐量引擎）服务于 LLMs。

抱歉，你的浏览器不支持播放此视频。

最困难的情况是模型太大以至于无法放入一个主机（假设一个张量并行切片跨 16 块芯片）。这时 Serve 只需一个额外字段 topology 就能解决问题。

了解这个字段是值得的，因为搞错它是典型的多主机 TPU 故障。当拓扑设置好后，Serve 的 TPU 后端会跳过通常的前置放置组，而是延迟到副本，由副本在启动时创建切片放置组。这种延迟使得张量并行模型的工作节点保持在一个共享的 ICI 网格上。若关闭该字段，Serve 将回退到按芯片的捆绑；在多主机模型中，这些捆绑可能会分散到两个切片上，而由于切片之间没有 ICI，工作节点无法完成第一次聚合操作。你不会遇到崩溃，而是会得到一个一直处于 DEPLOYING 状态的部署，同时你白白消耗 TPU 小时寻找实际上只缺一行 YAML 的错误。所以记住，topology 字段决定了差别。

在实际操作中，你会在一个发布的 vLLM TPU 镜像上部署 RayService（推荐用于生产环境而非原始 RayCluster），等待它达到 Running 状态，然后 curl 该端点。官方 GKE 教程涵盖 v5e 上的 Llama 3 8B 和 Mistral 7B，v6e 上的 Llama 3.1 70B，以及 Stable Diffusion。get-started 示例的 serve 步骤：https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started/serve 讲解了完整的端到端部署流程。

一个快速的加速器的有用性取决于你能否持续地向其输送数据，而 TPU 足够快，以至于原始加载器会成为瓶颈。这就是 iter_jax_batches() 解决的问题。它为你提供已经是 JAX 数组且已经设备分片好的批次，因此训练输入管道或大批量推理任务可以直接从 Ray Data 管道中拉取，无需主机端的 NumPy 到 JAX 复制来阻塞步骤。

iter_jax_batches API 会为你处理设备分片，并且它通过明确选择丢弃（drop）、填充（pad）或抛出（raise）来处理不完整的最后一个批次（即不是批次大小整倍数的批次），而不会在运行三小时后出现形状错误。

你可以把它用作 JaxTrainer 任务的输入端，它本身对于在 TPU 切片上进行大数据集的离线批处理推理也同样有用。它最近已加入 Ray，get-started 示例的数据步骤使用它来进行数据集准备和批处理推理。

TPU 上的 Ray Train：使用 JaxTrainer 进行分布式训练

在 TPU 上，训练曾经是 Ray 的一个令人困惑的部分，因为需要处理拓扑结构，并且在代码中考虑切片形状。JaxTrainer 解决了这一问题。它将 Ray Train 的训练循环（检查点、容错、多切片扩展）引入到 JAX 中——这是 Google 的数组和自动求导库，也是 TPU 的原生框架。你只需提供一个训练函数和一个切片形状，Ray 就会为每个主机启动一个工作器，将它们连接到单个网格中，并在每个工作器上运行你的函数。

在这个片段中，有两点你需要记住以节省调试时间。import jax 放在 train_loop_per_worker 内部，而不是文件顶部，因为每个工作器在其自己的 TPU 上下文中初始化 JAX；如果在模块作用域中导入，你将在第一步之前遇到神秘的设备初始化错误。topology="4x4" 就是整个布局声明，它取代了以前那一块手写的协调代码。与 GPU 上的 JaxTrainer 或 TorchTrainer 相比，唯一真正的区别是 use_tpu=True 和一个拓扑，而不是 GPU 数量。

其余部分就是运行。这是因为 Ray Train 拥有训练循环，你可以获得检查点和容错重启，这使得在可抢占容量上运行长时间的 TPU 训练实际上能够完成，并且当一个切片不足时，topology 可以扩展到多切片（Ray 会处理跨切片协调）。示例 get-started 的训练步骤：https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started/train 是一个完整的 JaxTrainer DPO 运行示例。

作为一流加速器支持的一部分，Ray 现在发布官方 rayproject/ray:*-tpu 镜像，内置 JAX/TPU 堆栈（jax[tpu]、flax、optax、orbax-checkpoint）和分析工具，因此你不必手动组装可用的 TPU 环境。你只需基于标记为 -tpu 的镜像即可。

在监控方面，Ray Dashboard——Ray 内置的集群和作业状态 Web UI——现在在 Cluster 标签页上显示 TPU 利用率和内存，与 CPU 和 GPU 一起显示，同时 ray.util.tpu.init_jax_profiler() 会暴露每个工作器的 JAX 分析器，仪表板可以附加使用。

在本开发者指南中，我们介绍了从 Ray 在 TPU 上的运行到 AI 工作负载运行的整个过程。

第1部分：https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ 显示，在TPU上运行Ray归结为一个注意事项，即将多主机模型保存在单个完整切片上，而GKE（通过Ray Operator插件）和Ray Core（通过slice_placement_group()）会为你处理这个问题。这部分在其之上放置了AI库：Ray Serve 利用单个 accelerator_config.topology 字段将多主机模型统一调度到一个切片上，Ray Data 通过 iter_jax_batches() 将切片传入 JAX 原生批次，JaxTrainer 则从一个 ScalingConfig 运行分布式训练循环。你已经在GPU上使用的同一个Ray，现在可以在TPU上使用。

更多内容即将到来。Google Cloud的Ray团队正在扩大TPU支持，详情见：https://discuss.google.dev/t/google-cloud-tpus-are-now-a-first-class-accelerator-in-ray/345281#p-914268-whats-next-in-progress-and-future-work-3：更深入的Ray Data和Ray LLM TPU集成，SkyRL 在多主机TPU上的强化学习和后训练，以及动态超/子切片支持都在规划中。对于你下一步，我的建议是：克隆入门示例：https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started，启动集群，然后运行 serve、data 或 train。或者只需在集群上启用 --enable-ray-operator，然后在小切片上运行一个Ray任务看看效果。使用TPU不必成为专家，只要尝试一下即可。现在，非常感谢你的阅读！如果你有任何额外问题或反馈，也可以在社交平台上联系我（LinkedIn：https://www.linkedin.com/in/ivan-nardini, X：https://x.com/ivnardini）。

新来的？第1部分：https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ 解释了切片、GKE以及Ray Core，这些是上面所有内容的基础。

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

弥合领域差距：使用Antigravity和Gemini构建的AI赛跑教练

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

TL;DR : Part 2 of 2. Part 1 ：https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ covered the one hardware idea you need and the two layers underneath (GKE and Ray Core). This part shows the libraries you actually build with, Ray Serve, Ray Data, and Ray Train.

Quick recap if you're landing here first. Running Ray on TPU comes down to one caveat: TPU chips are wired into fixed groups called slices (host VMs sharing a high-speed link called the ICI), and a multi-host model has to land on one whole slice or its workers can't reach each other and the job hangs.

Google Kubernetes Engine (GKE) with the Ray Operator add-on provisions slices and labels their hosts, and a Ray Core primitive, slice_placement_group() , reserves a whole slice at once. You declare a topology (the slice shape, like 4x4 for 16 chips, for example) and the libraries below handle the placement for you.

With Core handling placement underneath, the libraries all follow the same pattern: declare a topology, let Core reserve the slice. What changes per library is only what you declare it on. We'll go in the order most teams adopt them, serving first.

Serving is where most teams start from. A model that needs several GPUs to fit can run on a single TPU host, and TPUs are often a more available and cost-effective option for inference. Ray Serve gives you the usual autoscaling, load-balancing, and multi-model composition, and on TPU it serves LLMs through vLLM , a high-throughput engine.

Sorry, your browser doesn't support playback for this video

The hard case is when a model is too big for one host (say one sharded tensor-parallel across 16 chips). That's where Serve clears it with a single extra field, topology.

That one field is worth understanding, because getting it wrong is the classic multi-host TPU failure. With topology set, Serve's TPU backend skips its usual upfront placement group and defers to the replica, which creates a slice placement group at startup. That deferral is what keeps a tensor-parallel model's workers on one shared ICI mesh. Leave it off and Serve falls back to per-chip bundles; on a multi-host model those bundles can scatter across two slices, and because there's no ICI between slices, the workers never finish their first collective. You don't get a crash, you get a deployment that sits in DEPLOYING forever while you burn TPU-hours hunting for a bug that's really one missing line of YAML. So remember, topology field makes the difference.

In practice you deploy a RayService (recommended over a raw RayCluster for production) on a published vLLM TPU image, wait for it to reach Running, and curl the endpoint. The official GKE tutorials cover Llama 3 8B and Mistral 7B on v5e, Llama 3.1 70B on v6e, and Stable Diffusion. The serve step：https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started/serve of the get-started example walks the full deployment end to end.

A fast accelerator is only as useful as the data you can keep flowing into it, and TPUs are fast enough that a naive loader becomes the bottleneck. That's the problem iter_jax_batches() solves. It hands you batches that are already JAX arrays and already device-sharded, so a training input pipeline or a large batch-inference job pulls straight from a Ray Data pipeline with no host-side NumPy-to-JAX copy stalling the step.

iter_jax_batches API does the device sharding for you, and it handles the ragged final batch (the one that isn't a clean multiple of your batch size) with an explicit choice of drop, pad, or raise, instead of a shape error three hours into a run.

You can use it as the input side of a JaxTrainer job, and it's just as useful on its own for offline batch inference over a big dataset on a TPU slice. It landed recently in Ray, and the data step of the get-started example uses it for dataset prep and batch inference.

Ray Train on TPU: distributed training with JaxTrainer

Training used to be the confusing part of Ray on TPU, because of topology and having to account for the slice shape in your code. JaxTrainer addresses that. It brings Ray Train's training loop (checkpointing, fault tolerance, multi-slice scale-out) to JAX, Google's array and autodiff library and the native framework for TPU. You hand it a training function and a slice shape and Ray launches one worker per host, wires them into a single mesh, and runs your function on each.

Two things in this snippet you want to keep in mind to save debugging time. The import jax lives inside train_loop_per_worker , not at the top of the file, because each worker initializes JAX in its own TPU context; import it at module scope and you'll fight cryptic device-init errors before the first step. And topology="4x4" is the entire placement declaration, the line that used to be a block of hand-written coordination code. Set next to a GPU JaxTrainer or TorchTrainer, the only real difference is use_tpu=True and a topology instead of a GPU count.

The rest it just runs. This is because Ray Train owns the loop, you get checkpointing and fault-tolerant restarts, which is what makes long TPU runs on preemptible capacity actually finish, and topology scales to multi-slice (Ray wires the cross-slice coordination) when one slice isn't enough. The train step：https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started/train of the get-started example is a complete JaxTrainer DPO run.

As part of the first-class accelerator support, Ray now publishes official rayproject/ray:*-tpu images with the JAX/TPU stack ( jax[tpu] , flax, optax, orbax-checkpoint) and profiling tooling already installed, so you don't have to assemble a working TPU environment by hand. You can just base your image on the tagged -tpu one.

And for monitoring, the Ray Dashboard, Ray's built-in web UI for cluster and job state, now shows TPU utilization and memory next to CPU and GPU on the Cluster tab, with ray.util.tpu.init_jax_profiler() exposing a per-worker JAX profiler the dashboard can attach to.

In this developer guide on Ray on TPU, we covered the whole journey from how Ray runs on TPU to running AI workloads.

Part 1：https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ showed that running Ray on TPU comes down to one caveat, keeping a multi-host model on a single intact slice, and that GKE (through the Ray Operator add-on) and Ray Core (through slice_placement_group()) handle that for you. This part put the AI libraries on top: Ray Serve gang-schedules a multi-host model onto one slice with a single accelerator_config.topology field, Ray Data feeds the slice JAX-native batches through iter_jax_batches() , and JaxTrainer runs a distributed training loop from one ScalingConfig. The same Ray you already use on GPUs, now on TPU.

And more is coming. The Ray team on Google Cloud is widening TPU support from here：https://discuss.google.dev/t/google-cloud-tpus-are-now-a-first-class-accelerator-in-ray/345281#p-914268-whats-next-in-progress-and-future-work-3: deeper Ray Data and Ray LLM TPU integration, SkyRL on multi-host TPU for reinforcement learning and post-training, and dynamic super/sub-slice support are all on the roadmap. For your own next step, my recommendation is: clone the get-started example：https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started, stand up the cluster, then run serve, data, or train. Or just enable --enable-ray-operator on a cluster and run one Ray task on a small slice to see it work. You don't have to become a TPU expert to use one, just give it a try. For now, thanks for reading! And if you have any additional questions or feedback, feel free to reach out on socials (LinkedIn：https://www.linkedin.com/in/ivan-nardini, X：https://x.com/ivnardini).

New here? Part 1 ：https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ explains slices, GKE, and Ray Core, the foundation everything above builds on.

Expanding Choice in Gemini Enterprise Agent Platform: Introducing Grounding with Parallel Web Search

Bridging the Domain Gap: AI Race Coach built with Antigravity and Gemini

Run Ray on TPU, Part 1: The foundations