Ray now runs native PyTorch on Google TPUs, Anyscale and Google Cloud say

One worker per TPU chip

According to Anyscale, running TPU workloads on Ray used to mean adapting to a Single-Program, Multiple Data (SPMD) execution model designed primarily for JAX. In that model, one Ray worker process per host VM managed all of the machine's attached TPU chips.

Following TorchTPU, which the authors describe as Google's native PyTorch backend for Tensor Processing Units introduced last year, Ray now runs one worker per TPU chip. This matches the model PyTorch developers know from GPUs. Each worker requests a single logical TPU device with resources={"TPU": 1}, just as a GPU worker requests num_gpus=1.

• Device isolation: Ray assigns individual TPU devices to workers using TPU_VISIBLE_CHIPS, much as it sets CUDA_VISIBLE_DEVICES for GPU workers.
• Native device and collectives: migrating existing PyTorch code inside a Ray task or actor means changing "cuda" to "tpu" and switching the 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). Ray's SlicePlacementGroup API guarantees that actors are gang-scheduled together on the same TPU slice.

Training with TorchTrainer

For distributed training and fine-tuning, Ray Train's TorchTrainer can now run existing scripts on single-host and multi-host Cloud TPU slices with zero changes to the PyTorch training loop, according to the announcement. It supports Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDPv2) and DTensor workloads.

Teams configure TPUs through ScalingConfig by setting use_tpu=True, a target topology such as "4x4" and an accelerator_type such as "TPU-V6E" or "TPU-V7X". Ray Train then handles the orchestration:

• It provisions a SlicePlacementGroup so that all workers land on contiguous physical TPU slices over ICI, which Anyscale says prevents resource fragmentation.
• It spawns one worker per TPU device, sets the PyTorch distributed environment variables, and initializes the tpu_dist process group.
• Helper functions such as ray.train.torch.get_device() and ray.train.torch.prepare_model() detect the TorchTPU runtime, move the model to the TPU device and wrap it in DistributedDataParallel.

The blog post includes an example that scales a training loop across a 16-chip multi-host Trillium slice (TPU-V6E, 4x4 topology).

Inference with vLLM-TorchTPU and Ray Serve

For online inference, vLLM-TorchTPU registers TorchTPU as a vLLM platform plugin. PyTorch models can then compile and run on Cloud TPUs without rewriting model definitions in JAX, the authors say.

Combined with Ray Serve LLM, teams can deploy OpenAI-compatible vLLM endpoints across single-host and multi-host TPU slices using the same declarative LLMConfig API used on GPUs. When the accelerator settings specify a TPU slice, Ray Serve reserves a topology-aware SlicePlacementGroup for each replica and launches vLLM engine workers across it. The example in the post serves google/gemma-4-26B-A4B-it on a 4x4 TPU-V6E slice with a tensor parallel size of 16.

What it means for platform teams

Anyscale frames the release as an extension of the open-source PARK stack: PyTorch, AI models, Ray and Kubernetes. According to the company, PyTorch developers keep the eager execution model and Ray APIs they already know, while platform teams get GKE's topology-aware autoscaling, resilience and multi-accelerator governance.

Teams on GKE can now schedule, scale and operate distributed PyTorch training, vLLM inference and reinforcement learning workloads across GPUs and TPUs on a single Kubernetes control plane. Anyscale points to end-to-end GKE tutorials for serving LLMs and training models with Ray and TorchTPU. It also says it will expand the foundation to large-scale RL post-training frameworks and to next-generation Ironwood (TPU v7x) slices on GKE.