Run Ray on TPU, Half 1: The foundations

TL;DR: Should you already scale Python with Ray on GPUs, your code can now run on TPU (Tensor Processing Unit, Google’s AI accelerator chip) with totally supported official APIs you already know. The duty-and-actor mannequin, a JaxTrainer, the identical Ray Serve deployment simply pointed at TPUs orchestrated by Google Kubernetes Engine (GKE).

As of Ray 2.55, Google Cloud TPUs are a first-class accelerator in Ray. Which means that TPUs at the moment are in Ray’s launch pipelines with official pre-built pictures and help throughout the core libraries, as an alternative of the outdated “experimental” path the place you constructed your personal containers and leaned on group assist. On this “Run Ray on TPU” collection, you’ll find out how a TPU slice is simply one other accelerator Ray schedules onto (Half 1), then stroll by every library (Half 2).

Ray and TPU: A fast introduction

Ray is a distributed-computing framework: you write Python, and Ray runs it throughout a cluster as duties (stateless features) and actors (stateful staff). To Ray, a TPU is simply one other schedulable useful resource, the best way a GPU is. You ask for it, Ray locations your work on it.

However there may be one factor to bear in mind, after which we transfer on.

TPU chips are wired collectively into a set group referred to as a slice: a number of host machines (VMs) whose chips share a devoted high-speed hyperlink referred to as the ICI (Inter-Chip Interconnect). A multi-host mannequin has to land on one complete slice, or its staff cannot attain one another and the job simply hangs.

Should you assume in GPUs, image a slice as a single multi-GPU field the place the quick interconnect (NVLink) solely exists inside the field. Break up your staff throughout two bins with no cable between them and the collective operations, the all-reduce steps that synchronize gradients, by no means end. Coaching simply hangs. A TPU slice behaves the identical means: the ICI is that cable, and it solely reaches the chips of 1 slice.

That is the entire cause “Ray on TPU” wants something particular. One thing has to ensure all of your staff land on one intact slice. On GPUs you barely give it some thought; on TPUs it is essential, and Ray and GKE (Google Kubernetes Engine, Google’s managed Kubernetes) deal with it for you.

Yet one more phrase you may see in every single place is topology: it is the form of a slice, written like 4×4 for a 16-chip slice. You ask for a topology, not a chip depend.

When you perceive slicing and topology with TPU, the present Ray stack and your growth course of stays unchanged and it runs on TPU slices that GKE provisions. The diagram under is the entire system in a single image: on the left, the code you write (the Ray libraries you already use); within the center, the Ray Core layer that reserves complete slices; on the suitable, the GKE managed layer that provisions the {hardware} and labels it so Ray can discover slice boundaries.

GKE provisions a slice and labels its hosts, Ray Core reads these labels to order the entire slice directly, and your library name sits on prime, declaring a topology and nothing extra. No hand-written placement code wherever. The remainder of this half walks the underside two layers, GKE then Ray Core whereas Half 2 covers the Ray AI libraries.

How GKE orchestrates Ray on TPU

You run Ray on TPU by GKE utilizing the Ray Operator add-on.

# Autopilot (totally managed nodes)
gcloud container clusters create-auto CLUSTER 
  --enable-ray-operator --location=LOCATION

# or Customary (you handle node swimming pools)
gcloud container clusters create CLUSTER 
  --addons=RayOperator --location=LOCATION

Shell

That single flag installs two issues that matter for TPU. The primary, KubeRay, is the Kubernetes operator that turns RayCluster, RayService, and RayJob YAML into operating Ray clusters; it is the identical KubeRay you’d use with GPUs. The second is the TPU-specific half: the Ray TPU webhook, which stamps each TPU host with labels like ray.io/tpu-slice-name so Ray can inform which machines are wired into the identical slice. That label is the thread the entire system pulls on.

From there, you ask for TPUs in a manifest the identical means you’d ask for any node, with a nodeSelector for the technology and topology and the chip depend as a useful resource. A multi-host slice provides one discipline, numOfHosts.

# inside a RayCluster workerGroupSpec
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice   # the TPU technology
  cloud.google.com/gke-tpu-topology: "4x4"              # the slice form
# ... and request chips through the google.com/tpu useful resource restrict
numOfHosts: 4   # multi-host: what number of host VMs make up this slice

Shell

GKE provisions the slice, the webhook labels it, Ray reads the labels. You write Python. As soon as the add-on is up, you may see the KubeRay operator pod operating, and making use of that manifest brings up a head pod plus one employee pod per host within the slice. The cluster step of the get-started instance provisions all of this with Terraform.

What really retains your staff collectively is a Ray Core primitive sitting simply above this layer, the slice placement group, and that is the place the remainder of the information begins.

Ray Core on TPU

Ray Core is the bottom layer, the task-and-actor engine and scheduler every thing else sits on. Its TPU help lives within the public ray.util.tpu API, and there is actually one operate to know: slice_placement_group(). It takes that “preserve my staff on one intact slice” assure from earlier and turns it right into a single name, reserving an entire slice atomically (all hosts or none) by matching on the webhook labels.

from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

# Reserve one complete v6e 4x4 slice (16 chips throughout 4 hosts), atomically
spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.prepared(), timeout=600)

@ray.distant(sources={"TPU": 4})
def employee(rank, world): ...

duties = [
    worker.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=spg.placement_group)
    ).remote(rank=i, world=spg.num_hosts)
    for i in range(spg.num_hosts)
]

Python

It is very important spotlight that you not often name slice_placement_group your self. The Ray AI libraries (Information, Practice, Serve) name it for you, so in observe you declare a topology they usually deal with the slice. You’d solely attain for slice_placement_group() immediately once you’re writing a customized distributed workload that is not Practice, Serve, or Information. One caveat price understanding: the API is public however marked alpha (@PublicAPI(stability="alpha")), so it is usable immediately however the floor can nonetheless shift between releases.

That is the inspiration. Subsequent are the libraries.

You now have the entire psychological mannequin: a slice has to remain intact, GKE provisions and labels it, and Ray Core reserves it as a unit so that you by no means hand-write placement code. The whole lot you really construct sits on prime of that and reuses it.

In Half 2, we are going to discover how you should use Ray AI libraries on TPU for serving LLMs with vLLM, feeding slices with Ray Information and coaching with JaxTrainer.