Announcing Ray support for PyTorch on Google TPUs

Earlier this year at Ray Summit 2026, we covered how 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 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, 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 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:

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:

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.

LinkvLLM-TorchTPU with Ray Serve

For online inference, 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, 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.

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 or Training models 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.