{"slug": "run-ray-on-tpu-part-2-ray-ai-libraries", "title": "Run Ray on TPU, Part 2: Ray AI libraries", "summary": "Ray AI libraries (Serve, Data, Train) now support Google TPU slices through a topology field that reserves a whole ICI-connected slice, preventing multi-host deployment hangs. Ray Serve serves LLMs via vLLM on TPUs with autoscaling and load-balancing, while Ray Data's iter_jax_batches() delivers device-sharded JAX arrays to eliminate host-side copy bottlenecks. The official GKE tutorials cover Llama 3 8B and Mistral 7B on v5e, Llama 3.1 70B on v6e, and Stable Diffusion.", "body_md": "*TL;DR**: Part 2 of 2.* *Part 1**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.*\n\n**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.\n\nGoogle Kubernetes Engine (GKE) with the Ray Operator add-on provisions slices and labels their hosts, and a Ray Core primitive, `slice_placement_group()`\n\n, 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.\n\nWith 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.\n\nServing 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.\n\nThe 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.\n\n```\naccelerator_type: TPU-V6E\naccelerator_config:\n  kind: tpu\n  topology: \"4x4\"\n```\n\nThat 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.\n\nIn 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.\n\nA 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()`\n\nsolves. 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.\n\n```\nds = ray.data.read_parquet(\"gs://my-bucket/train/\")\nfor batch in ds.iter_jax_batches(batch_size=1024):\n    # batch arrives as device-sharded JAX arrays, ready for the training step\n    loss = train_step(batch)\n```\n\n`iter_jax_batches`\n\nAPI 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.\n\nYou can use it as the input side of a `JaxTrainer`\n\njob, 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.\n\nRay Train on TPU: distributed training with JaxTrainer\n\nTraining 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`\n\naddresses 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.\n\n``` python\nfrom ray.train import ScalingConfig\nfrom ray.train.v2.jax import JaxTrainer\n\ndef train_loop_per_worker(config):\n    import jax            # import jax INSIDE the worker fn (TPU requirement)\n    # ... your JAX/Flax training step runs here, once per host ...\n\ntrainer = JaxTrainer(\n    train_loop_per_worker=train_loop_per_worker,\n    scaling_config=ScalingConfig(\n        use_tpu=True,\n        topology=\"4x4\",            # the slice shape, NOT a chip count\n        accelerator_type=\"TPU-V6E\",\n    ),\n)\ntrainer.fit()\n```\n\nTwo things in this snippet you want to keep in mind to save debugging time. The import jax lives *inside* `train_loop_per_worker`\n\n, 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\"`\n\nis 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.\n\nThe 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.\n\nAs part of the first-class accelerator support, Ray now publishes official `rayproject/ray:*-tpu`\n\nimages with the JAX/TPU stack (`jax[tpu]`\n\n, 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`\n\none.\n\nAnd 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()`\n\nexposing a per-worker JAX profiler the dashboard can attach to.\n\nIn this developer guide on Ray on TPU, we covered the whole journey from how Ray runs on TPU to running AI workloads.\n\n[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())`\n\nhandle that for you.\n\nThis part put the AI libraries on top: Ray Serve gang-schedules a multi-host model onto one slice with a single `accelerator_config.topology`\n\nfield, Ray Data feeds the slice JAX-native batches through `iter_jax_batches()`\n\n, and JaxTrainer runs a distributed training loop from one ScalingConfig. The same Ray you already use on GPUs, now on TPU.\n\nAnd 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.\n\nFor 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`\n\non 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.\n\nFor 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)).\n\nHappy building!\n\n`kubernetes-engine-samples`\n\n, the serve, data, and train steps from this post as working code (Qwen3-4B on a v6e slice).*New here?* *Part 1**explains slices, GKE, and Ray Core, the foundation everything above builds on.*", "url": "https://wpnews.pro/news/run-ray-on-tpu-part-2-ray-ai-libraries", "canonical_source": "https://developers.googleblog.com/run-ray-on-tpu-part-2-ray-ai-libraries/", "published_at": "2026-07-24 16:09:12.897391+00:00", "updated_at": "2026-07-24 16:09:15.538874+00:00", "lang": "en", "topics": ["artificial-intelligence", "ai-infrastructure", "ai-tools", "large-language-models", "machine-learning"], "entities": ["Ray", "Google TPU", "Google Kubernetes Engine", "vLLM", "Llama 3", "Mistral 7B", "Llama 3.1 70B", "Stable Diffusion"], "alternates": {"html": "https://wpnews.pro/news/run-ray-on-tpu-part-2-ray-ai-libraries", "markdown": "https://wpnews.pro/news/run-ray-on-tpu-part-2-ray-ai-libraries.md", "text": "https://wpnews.pro/news/run-ray-on-tpu-part-2-ray-ai-libraries.txt", "jsonld": "https://wpnews.pro/news/run-ray-on-tpu-part-2-ray-ai-libraries.jsonld"}}