# Announcing Ray support for PyTorch on Google TPUs

> Source: <https://anyscale.com/blog/announcing-ray-support-native-pytorch-google-tpu>
> Published: 2026-10-20 09:00:00+00:00

# 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.