Loading market data...

Ray Adds Native PyTorch Support for Google TPUs on GKE

Ray Adds Native PyTorch Support for Google TPUs on GKE

Anyscale and Google Cloud said Tuesday that Ray now supports PyTorch on Google TPUs natively on Google Kubernetes Engine (GKE). The integration lets teams run PyTorch workloads on TPUs through Ray without custom glue code, and it uses one Ray worker per TPU chip.

The companies also built in TPU-aware scheduling, so Ray places tasks with the underlying TPU topology in mind instead of treating the accelerators as generic compute. That matters for training and tuning jobs where a misplaced worker can leave a chip idle or force a costly reshuffle.

What the native integration changes

Until now, running PyTorch on TPUs under Ray meant stitching together environment setup, device assignment and scheduling logic. The new support moves those steps into Ray itself on GKE. A Ray worker maps to a single TPU chip, which gives the scheduler a clear unit to place and the runtime a clear unit to track.

For teams already using Ray for distributed training, the change is mostly about removing friction. Jobs that previously needed custom launch scripts can target TPUs the same way they target other accelerators. The TPU-aware scheduler is the piece that makes the one-worker-per-chip model practical at scale, because it can account for how chips are grouped on a host rather than scattering workers across the cluster.

Why GKE is the delivery vehicle

Google Cloud is positioning GKE as the place to run these workloads. Kubernetes already handles provisioning and isolation for TPU node pools, and Ray sits on top to manage the distributed application. The combination is aimed at teams that want to keep orchestration in Kubernetes while using Ray for the Python-level distribution.

The announcement doesn't include pricing, performance benchmarks or a list of supported TPU generations. Those details will matter to anyone deciding whether to move existing GPU-based Ray jobs onto TPUs. For now, the claim is narrower: native support exists, and it runs on GKE.

Who this is for

The obvious audience is Ray users who have wanted TPU capacity without maintaining a separate stack. PyTorch is the default framework for a large share of those teams, and TPUs are Google's own silicon, so the integration closes a gap that had been filled by community workarounds and internal tooling.

It also gives Google Cloud a way to make TPUs more accessible to the Ray ecosystem without asking users to abandon Kubernetes. Anyscale, for its part, gets another supported accelerator path for its managed Ray offerings. Neither company disclosed how long the integration was in development or which customers have tested it.

What's still missing

There's no word yet on when the feature will reach general availability, which TPU versions are covered, or whether the same native path will extend beyond GKE to other Kubernetes distributions. Teams evaluating a move will also want numbers on throughput and cost per chip compared with their current setup.

For now, the integration is available to try on GKE. The next concrete step for interested teams is to check whether their Ray version and TPU node pools meet the prerequisites — and to watch for the benchmark data that usually follows an announcement like this.