TPU で Ray を動かす Part 1: 基礎編Run Ray on TPU, Part 1: The foundations

Google has published the first installment of a technical series explaining how to run Ray, the popular open-source distributed computing framework, on its Tensor Processing Units (TPUs). The pairing matters because Ray has become a common orchestration layer for large-scale machine learning, while TPUs are among the most powerful accelerators available for training and serving large models, and bridging the two has historically required careful configuration.

Ray originated at UC Berkeley's RISELab and is now developed largely by Anyscale. At its core, Ray provides a general-purpose model for distributed execution built around two primitives: stateless tasks and stateful actors. On top of Ray Core sit higher-level libraries such as Ray Train for distributed model training, Ray Serve for inference, Ray Data for distributed data processing, and Ray Tune for hyperparameter search. This layered design lets teams scale Python workloads across a cluster without rewriting their application logic, which is one reason Ray has been adopted across the industry for orchestrating training pipelines and reinforcement learning systems.

TPUs, by contrast, are Google's custom application-specific integrated circuits designed for the linear algebra that dominates neural network computation. They are typically programmed through XLA-compiled frameworks such as JAX, TensorFlow, or PyTorch/XLA. A key architectural detail is that TPUs are organized into pods and slices, where multiple chips are connected by a high-speed interconnect. Larger configurations span many host machines, and each host is attached to a set of TPU chips. This multi-host topology is central to understanding how Ray fits in, because a distributed framework must map its own notion of nodes and workers onto the physical arrangement of TPU hosts and chips.

The foundational piece appears to focus on that mapping. In a typical setup, a Ray cluster consists of a head node that manages cluster state and scheduling, plus a number of worker nodes that execute tasks and actors. When running on TPUs, each TPU host generally corresponds to a Ray worker, and the accelerators attached to that host become resources that Ray can schedule against. Google's guidance covers how to represent TPUs as custom resources so that Ray's scheduler places workloads on the correct hosts, an important consideration given that many TPU programs use a single-program, multiple-data (SPMD) execution model that expects all participating hosts to run in lockstep.

That difference in execution philosophy is part of why the integration is worth documenting. Ray's task and actor model is flexible and dynamic, whereas large TPU training jobs often assume a tightly coordinated, gang-scheduled group of processes. Reconciling the two means ensuring that Ray launches one worker per host, that those workers can discover one another, and that the collective communication libraries underlying JAX or PyTorch/XLA can establish the necessary connections across the slice. The series' emphasis on design philosophy suggests it aims to give readers a mental model for these tradeoffs rather than only a set of commands.

For context, this work sits alongside several related Google efforts. TPU VMs allow users to run code directly on the machines attached to the accelerators, which simplifies frameworks like Ray that assume direct host access. Google also offers KubeRay, an operator for running Ray on Kubernetes, and Ray on TPUs can be deployed through Google Kubernetes Engine, giving teams a managed path to provisioning TPU slices. Separately, Google has promoted Pathways, its own orchestration system for large-scale accelerator workloads, so Ray represents an open-source alternative that many organizations already know.

The broader significance is that combining Ray with TPUs lowers a barrier for teams that have standardized on Ray but want access to Google's accelerators without abandoning their existing tooling. It also reflects a wider industry trend in which distributed frameworks increasingly need first-class support for heterogeneous hardware, from GPUs to custom silicon. As a Part 1, the article is explicitly foundational, and readers should expect subsequent installments to build on these basics with more advanced topics such as multi-slice scaling, fault tolerance, and end-to-end training or serving examples. For now, it establishes the vocabulary and architecture needed to understand how a general distributed system and a specialized accelerator can work together.