[Core][Train] Support Torch runtime with TPUs
Author: ryanaolearyCreated Jun 12, 2026Updated Sep 17, 2026
Labelsenhancementtriagetrainperformancecommunity-backlog
Description
This can serve as a tracking issue for efforts related to enabling support for PyTorch on TPUs with Ray.
- Inject Torch TPU required vars to RayCluster Pods requesting TPU
- Fix resource accounting for workloads that require per-device processes
- Update
TorchTrainerto support TPU backend and add validating tests - Implement multi-slice support in Ray Train's TorchTrainer
- Support MPMD workloads with TorchTPU and Ray
- Fix
TPU_VISIBLE_CHIPSresource bounding for privileged K8s containers on GKE - Support TorchTPU with
SubSlicePlacementGroups - Add physical topology ordering to
ray.io/tpu-worker-idused by Ray and Torch/vLLM frameworks
- Fix
Use case
No response
Source: ray-project/ray