{"slug": "announcing-ray-support-for-pytorch-on-google-tpus", "title": "Announcing Ray support for PyTorch on Google TPUs", "summary": "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.", "body_md": "# Announcing Ray support for PyTorch on Google TPUs\n\n[Lehui Liu](https://anyscale.com/blog?author=lehui-liu),\n\n[Richard Liaw](https://anyscale.com/blog?author=richard-liaw),\n\n[Andrew Sy Kim (Google)](https://anyscale.com/blog?author=andrew-sy-kim-google),\n\n[Nisha Mariam Johnson](https://anyscale.com/blog?author=nisha-mariam-johnson)and\n\n[Ryan O'Leary (Google)](https://anyscale.com/blog?author=ryan-oleary-google)| October 20, 2026\n\nEarlier 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.\n\nToday 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.\n\nBuilding 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.\n\n## LinkNative PyTorch on Google TPUs\n\nHistorically, 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. \n\nFollowing 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.\n\nWhen 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`:\n\n- **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`\n- **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\"` .\n- **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.\n\nBelow 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:\n\n``` python\nimport os\nimport ray\nimport torch\nimport torch.distributed as dist\n\n# On GPU: @ray.remote(num_gpus=1)\n@ray.remote(resources={\"TPU\": 1})\nclass TorchWorker:\n    def __init__(self, rank: int, world_size: int):\n        self.rank = rank\n        self.world_size = world_size\n        # Ray automatically isolates the assigned device via TPU_VISIBLE_CHIPS\n        self.local_rank = os.environ[\"TPU_VISIBLE_CHIPS\"]\n\n    def run_all_reduce(self, master_addr: str, master_port: int):\n        os.environ[\"MASTER_ADDR\"] = master_addr\n        os.environ[\"MASTER_PORT\"] = str(master_port)\n        os.environ[\"WORLD_SIZE\"] = str(self.world_size)\n        os.environ[\"RANK\"] = str(self.rank)\n        os.environ[\"LOCAL_RANK\"] = self.local_rank\n\n        # On GPU: dist.init_process_group(backend=\"nccl\")\n        dist.init_process_group(backend=\"tpu_dist\")\n\n        # On GPU: device = torch.device(\"cuda\")\n        device = torch.device(\"tpu\")\n        tensor = torch.ones(2, 2, device=device) * self.rank\n\n        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)\n        return tensor.cpu().numpy()\n```\n\n## LinkTorchTrainer on TPUs\n\nFor 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. \n\nWith 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.\n\nWhen you configure `TorchTrainer` for TPUs via `ScalingConfig`, Ray Train handles all of the underlying TPU orchestration automatically:\n\n1. **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.\n2. **Automated** **tpu_dist**\n **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.\n3. **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.\n\nHere is a complete example of scaling a PyTorch training loop across a 16-chip multi-host Trillium (`TPU-V6E`, `4x4` topology) slice using `TorchTrainer`:\n\n``` python\nimport ray\nfrom ray import train\nfrom ray.train import ScalingConfig\nfrom ray.train.torch import TorchTrainer\nimport torch\nfrom torch import nn, optim\nimport torch.distributed as dist\n\ndef train_fn_per_worker(config):\n    # Automatically resolves to torch.device(\"tpu\") on TPU workers\n    device = ray.train.torch.get_device()\n\n    model = nn.Sequential(\n        nn.Linear(10, 32),\n        nn.ReLU(),\n        nn.Linear(32, 1),\n    ).to(device)\n\n    # Wraps the model in DDP over the \"tpu_dist\" backend\n    model = ray.train.torch.prepare_model(model)\n\n    criterion = nn.MSELoss()\n    optimizer = optim.SGD(model.parameters(), lr=0.05)\n\n    inputs = torch.randn(128, 10, device=device)\n    targets = inputs.sum(dim=1, keepdim=True)\n\n    for epoch in range(10):\n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss = criterion(outputs, targets)\n        loss.backward()\n        optimizer.step()\n\n        train.report({\"loss\": loss.item(), \"epoch\": epoch + 1})\n\nscaling_config = ScalingConfig(\n    num_workers=16,\n    use_tpu=True,\n    topology=\"4x4\",\n    accelerator_type=\"TPU-V6E\",\n    resources_per_worker={\"TPU\": 1},\n)\n\ntrainer = TorchTrainer(\n    train_loop_per_worker=train_fn_per_worker,\n    scaling_config=scaling_config,\n)\n\nresult = trainer.fit()\n```\n\nFor 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).\n\n## LinkvLLM-TorchTPU with Ray Serve\n\nFor 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.\n\nWhen `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:\n\n```\nfrom ray.serve.llm import (\n    LLMConfig,\n    LLMServingArgs,\n    ModelLoadingConfig,\n    build_openai_app,\n)\nllm_config = LLMConfig(\n    model_loading_config=ModelLoadingConfig(\n        model_id=\"google/gemma-4-26B-A4B-it\",\n    ),\n    # On GPU: accelerator_type=\"H100\"\n    accelerator_type=\"TPU-V6E\",\n    accelerator_config={\n        \"kind\": \"tpu\",\n        \"topology\": \"4x4\",\n    },\n    placement_group_config={\n        \"bundle_per_worker\": {\"TPU\": 1},\n    },\n    engine_kwargs={\n        \"tensor_parallel_size\": 16,\n        \"enable_expert_parallel\": True,\n        \"max_model_len\": 8192,\n    },\n)\napp = build_openai_app(LLMServingArgs(llm_configs=[llm_config]))\n```\n\nFor 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).\n\n## LinkExtending the PARK stack to TPUs\n\nAs 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.\n\nTo 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.", "url": "https://wpnews.pro/news/announcing-ray-support-for-pytorch-on-google-tpus", "canonical_source": "https://anyscale.com/blog/announcing-ray-support-native-pytorch-google-tpu", "published_at": "2026-10-20 09:00:00+00:00", "updated_at": "2026-10-20 15:16:20.864004+00:00", "lang": "en", "topics": ["ai-infrastructure", "machine-learning", "mlops", "ai-chips", "developer-tools"], "entities": ["Anyscale", "Google Cloud", "Ray", "PyTorch", "Google TPUs", "TorchTPU", "KubeRay", "Google Kubernetes Engine"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/announcing-ray-support-for-pytorch-on-google-tpus", "markdown": "https://wpnews.pro/news/announcing-ray-support-for-pytorch-on-google-tpus.md", "text": "https://wpnews.pro/news/announcing-ray-support-for-pytorch-on-google-tpus.txt", "jsonld": "https://wpnews.pro/news/announcing-ray-support-for-pytorch-on-google-tpus.jsonld"}}