Announcing Ray support for PyTorch on Google TPUs

Anyscale and Google Cloud announced native Ray support for PyTorch on Google TPUs on October 20, 2026, letting GKE users schedule distributed PyTorch training, vLLM inference and reinforcement learning workloads across GPUs and TPUs on a single Kubernetes control plane. The integration builds on Google's open-source TorchTPU project and KubeRay, and changes Ray's TPU model from one worker process per host VM managing all attached chips to one worker per TPU chip, with each worker requesting resources={"TPU": 1} and device isolation via TPU_VISIBLE_CHIPS. Migrating code from GPUs to TPUs now requires changing "cuda" to "tpu" and the communications backend from "nccl" to "tpu_dist", while Ray's SlicePlacementGroup API gang-schedules actors on the same multi-host TPU slice.

Announcing Ray support for PyTorch on Google TPUs Lehui Liu https://anyscale.com/blog?author=lehui-liu , Richard Liaw https://anyscale.com/blog?author=richard-liaw , Andrew Sy Kim Google https://anyscale.com/blog?author=andrew-sy-kim-google , Nisha Mariam Johnson https://anyscale.com/blog?author=nisha-mariam-johnson and Ryan O'Leary Google https://anyscale.com/blog?author=ryan-oleary-google | October 20, 2026 Earlier this year at Ray Summit 2026, we covered how the PARK stack https://www.oreilly.com/radar/what-is-the-park-stack/ PyTorch, AI models, Ray, and Kubernetes is emerging as an industry standard for production AI. Ray has long been the ubiquitous distributed runtime for PyTorch workloads on GPUs. Today Anyscale and Google Cloud are announcing native Ray support for PyTorch on Google TPUs, the foundational supercomputing infrastructure powering Gemini and large-scale AI workloads on Google Cloud. Building on Google's TorchTPU project https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus and KubeRay on Google Kubernetes Engine GKE , this integration brings a native, single-device-per-worker PyTorch experience to Cloud TPUs. Platform and ML teams running on GKE can now schedule, scale, and operate distributed PyTorch training, vLLM inference, and reinforcement learning RL workloads across GPUs and TPUs on a single Kubernetes control plane. LinkNative PyTorch on Google TPUs Historically, running TPU workloads on Ray required adapting to a Single-Program, Multiple Data SPMD execution model designed primarily for JAX, where a single Ray worker process per host VM managed all of the machine's attached TPU chips. TorchTPU https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus , introduced last year as Google's native PyTorch backend for Tensor Processing Units,brings an eager-first, native TPU integration and a torch.distributed extension for running distributed workloads on TPUs. Following TorchTPU, we’ve evolved Ray to run one worker per TPU chip, matching the single-device-per-worker model that PyTorch developers are used to on GPUs. When scheduling tasks or actors with Ray and TorchTPU, each worker requests a single logical TPU device resources={"TPU": 1} just as a GPU worker requests num gpus=1 : - Device Isolation: Just as Ray automatically sets CUDA VISIBLE DEVICES for GPU workers, Ray isolates and assigns individual TPU devices to workers using TPU VISIBLE CHIPS - Native PyTorch Device & Collectives: Inside a Ray task or actor, migrating existing PyTorch code from GPUs to TPUs is as simple as changing "cuda" to "tpu" and switching the distributed communications backend from "nccl" to "tpu dist" . - Topology-Aware Gang Scheduling: GKE provisions multi-host TPU slices as node pools connected over high-speed Inter-Chip Interconnects ICI . To take advantage of this physical layout, Ray’s SlicePlacementGroup https://docs.ray.io/en/latest/ray-core/api/doc/ray.util.tpu.SlicePlacementGroup.html API provides a native scheduling implementation that guarantees actors are gang-scheduled together on the same TPU slice. Below is an example of a distributed PyTorch all reduce using Ray actors on TPUs. Aside from the resource request and device/backend strings, the code is identical to a Ray actor running on GPUs: python import os import ray import torch import torch.distributed as dist On GPU: @ray.remote num gpus=1 @ray.remote resources={"TPU": 1} class TorchWorker: def init self, rank: int, world size: int : self.rank = rank self.world size = world size Ray automatically isolates the assigned device via TPU VISIBLE CHIPS self.local rank = os.environ "TPU VISIBLE CHIPS" def run all reduce self, master addr: str, master port: int : os.environ "MASTER ADDR" = master addr os.environ "MASTER PORT" = str master port os.environ "WORLD SIZE" = str self.world size os.environ "RANK" = str self.rank os.environ "LOCAL RANK" = self.local rank On GPU: dist.init process group backend="nccl" dist.init process group backend="tpu dist" On GPU: device = torch.device "cuda" device = torch.device "tpu" tensor = torch.ones 2, 2, device=device self.rank dist.all reduce tensor, op=dist.ReduceOp.SUM return tensor.cpu .numpy LinkTorchTrainer on TPUs For distributed training and fine-tuning, many teams rely on Ray Train’s TorchTrainer to orchestrate Distributed Data Parallel DDP , Fully Sharded Data Parallel FSDPv2 , and DTensor workloads without writing manual boilerplate for distributed process groups. With native TorchTPU support in Ray Train, you can now run your existing TorchTrainer scripts on single-host and multi-host Cloud TPU slices with zero changes to your PyTorch training loop. When you configure TorchTrainer for TPUs via ScalingConfig , Ray Train handles all of the underlying TPU orchestration automatically: 1. Atomic TPU Slice Reservation: By specifying use tpu=True alongside your target topology e.g., "4x4" and accelerator type e.g., "TPU-V6E" or "TPU-V7X" , TorchTrainer automatically provisions a SlicePlacementGroup under the hood. This guarantees that all training workers are gang-scheduled onto contiguous physical TPU slices over high-speed ICI, preventing resource fragmentation. 2. Automated tpu dist Bootstrap: Ray Train spawns one worker per TPU device resources per worker={"TPU": 1} , automatically configures PyTorch distributed environment variables RANK , WORLD SIZE , LOCAL RANK , MASTER ADDR , MASTER PORT alongside TorchTPU’s slice-builder addresses and mesh topology, and initializes the tpu dist process group. 3. Drop-in Ray Train Utilities: Standard Ray Train helper functions like ray.train.torch.get device and ray.train.torch.prepare model automatically detect the TorchTPU runtime, move your model to the assigned tpu device, and wrap it in PyTorch’s DistributedDataParallel class. Here is a complete example of scaling a PyTorch training loop across a 16-chip multi-host Trillium TPU-V6E , 4x4 topology slice using TorchTrainer : python import ray from ray import train from ray.train import ScalingConfig from ray.train.torch import TorchTrainer import torch from torch import nn, optim import torch.distributed as dist def train fn per worker config : Automatically resolves to torch.device "tpu" on TPU workers device = ray.train.torch.get device model = nn.Sequential nn.Linear 10, 32 , nn.ReLU , nn.Linear 32, 1 , .to device Wraps the model in DDP over the "tpu dist" backend model = ray.train.torch.prepare model model criterion = nn.MSELoss optimizer = optim.SGD model.parameters , lr=0.05 inputs = torch.randn 128, 10, device=device targets = inputs.sum dim=1, keepdim=True for epoch in range 10 : optimizer.zero grad outputs = model inputs loss = criterion outputs, targets loss.backward optimizer.step train.report {"loss": loss.item , "epoch": epoch + 1} scaling config = ScalingConfig num workers=16, use tpu=True, topology="4x4", accelerator type="TPU-V6E", resources per worker={"TPU": 1}, trainer = TorchTrainer train loop per worker=train fn per worker, scaling config=scaling config, result = trainer.fit For a complete end-to-end walkthrough on Google Kubernetes Engine GKE , see Train models on multi-host TPUs with Ray and TorchTPU https://docs.cloud.google.com/kubernetes-engine/docs/tutorials/train-models-tpu-ray-torchtpu . LinkvLLM-TorchTPU with Ray Serve For online inference, vLLM-TorchTPU https://github.com/vllm-project/vllm-torchtpu registers TorchTPU as a vLLM platform plugin so PyTorch models compile and execute on Cloud TPUs without rewriting model definitions in JAX. Combined with Ray Serve LLM https://docs.ray.io/en/latest/serve/llm/index.html , you can deploy OpenAI-compatible vLLM endpoints across single-host and multi-host TPU slices using the same declarative LLMConfig API used on GPUs. When accelerator type and accelerator config specify a TPU slice, Ray Serve reserves a topology-aware SlicePlacementGroup for each replica and launches vLLM engine workers across the slice: from ray.serve.llm import LLMConfig, LLMServingArgs, ModelLoadingConfig, build openai app, llm config = LLMConfig model loading config=ModelLoadingConfig model id="google/gemma-4-26B-A4B-it", , On GPU: accelerator type="H100" accelerator type="TPU-V6E", accelerator config={ "kind": "tpu", "topology": "4x4", }, placement group config={ "bundle per worker": {"TPU": 1}, }, engine kwargs={ "tensor parallel size": 16, "enable expert parallel": True, "max model len": 8192, }, app = build openai app LLMServingArgs llm configs= llm config For a complete end-to-end walkthrough on Google Kubernetes Engine GKE , see Serve LLMs on multi-host TPUs with Ray and TorchTPU https://cloud.google.com/kubernetes-engine/docs/tutorials/serve-llm-tpu-ray-torchtpu . LinkExtending the PARK stack to TPUs As AI workloads evolve from standalone training jobs to compound systems combining training, high-throughput inference, and reinforcement learning, organizations are standardizing on the open-source PARK stack PyTorch, AI models, Ray, and Kubernetes . Bringing TorchTPU to Ray and GKE completes this vision on Google Cloud: PyTorch developers keep the eager execution model and Ray APIs they already know, while platform teams get the topology-aware autoscaling, resilience, and multi-accelerator governance of GKE. To get started today, explore the end-to-end GKE tutorials for Serving LLMs https://cloud.google.com/kubernetes-engine/docs/tutorials/serve-llm-tpu-ray-torchtpu or Training models https://docs.cloud.google.com/kubernetes-engine/docs/tutorials/train-models-tpu-ray-torchtpu with Ray and TorchTPU, and stay tuned as we expand this foundation to large-scale RL post-training frameworks and next-generation Ironwood TPU v7x slices on GKE.