Ray unterstützt jetzt natives PyTorch auf Google TPUs, melden Anyscale und Google Cloud

Ein Worker pro TPU-Chip

Laut Anyscale musste man TPU-Workloads auf Ray bislang an ein Single-Program-Multiple-Data-Ausführungsmodell (SPMD) anpassen, das vor allem für JAX entworfen wurde. Dabei verwaltete ein einzelner Ray-Worker-Prozess pro Host-VM sämtliche angeschlossenen TPU-Chips der Maschine.

Nach dem Vorbild von TorchTPU, das die Autoren als Googles natives PyTorch-Backend für Tensor Processing Units beschreiben und das im vergangenen Jahr vorgestellt wurde, startet Ray nun einen Worker pro TPU-Chip. Das entspricht dem Modell, das PyTorch-Entwickler von GPUs kennen. Jeder Worker fordert mit resources={"TPU": 1} ein einzelnes logisches TPU-Gerät an, genau wie ein GPU-Worker num_gpus=1 anfordert.

• Geräteisolierung: Ray weist Workern einzelne TPU-Geräte über TPU_VISIBLE_CHIPS zu, ähnlich wie es für GPU-Worker CUDA_VISIBLE_DEVICES setzt.
• Natives Gerät und native Collectives: Um bestehenden PyTorch-Code in einem Ray-Task oder -Actor zu migrieren, genügt es, "cuda" durch "tpu" zu ersetzen und das Kommunikations-Backend von "nccl" auf "tpu_dist" umzustellen.
• Topologiebewusstes Gang Scheduling: GKE stellt Multi-Host-TPU-Slices als Node Pools bereit, die über schnelle Inter-Chip Interconnects (ICI) verbunden sind. Die SlicePlacementGroup-API von Ray garantiert, dass Actors gemeinsam auf demselben TPU-Slice eingeplant werden.

Training mit TorchTrainer

Für verteiltes Training und Fine-Tuning kann der TorchTrainer von Ray Train laut Ankündigung nun bestehende Skripte auf Single-Host- und Multi-Host-Cloud-TPU-Slices ausführen, ohne dass die PyTorch-Trainingsschleife geändert werden muss. Unterstützt werden Workloads mit Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDPv2) und DTensor.

Teams konfigurieren TPUs über ScalingConfig, indem sie use_tpu=True, eine Zieltopologie wie "4x4" und einen accelerator_type wie "TPU-V6E" oder "TPU-V7X" festlegen. Die Orchestrierung übernimmt dann Ray Train:

• Es richtet eine SlicePlacementGroup ein, damit alle Worker auf zusammenhängenden physischen TPU-Slices über ICI landen. Das verhindert laut Anyscale eine Fragmentierung der Ressourcen.
• Es startet pro TPU-Gerät einen Worker, setzt die verteilten PyTorch-Umgebungsvariablen und initialisiert die Prozessgruppe tpu_dist.
• Hilfsfunktionen wie ray.train.torch.get_device() und ray.train.torch.prepare_model() erkennen die TorchTPU-Laufzeitumgebung, verschieben das Modell auf das TPU-Gerät und verpacken es in DistributedDataParallel.

Der Blogbeitrag enthält ein Beispiel, in dem eine Trainingsschleife auf einem Multi-Host-Trillium-Slice mit 16 Chips (TPU-V6E, Topologie 4x4) skaliert wird.

Inferenz mit vLLM-TorchTPU und Ray Serve

Für Online-Inferenz registriert vLLM-TorchTPU TorchTPU als Plattform-Plugin für vLLM. PyTorch-Modelle lassen sich dadurch auf Cloud TPUs kompilieren und ausführen, ohne dass Modelldefinitionen in JAX neu geschrieben werden müssen, so die Autoren.

In Kombination mit Ray Serve LLM können Teams OpenAI-kompatible vLLM-Endpunkte auf Single-Host- und Multi-Host-TPU-Slices bereitstellen, und zwar mit derselben deklarativen LLMConfig-API, die auch auf GPUs zum Einsatz kommt. Legen die Beschleunigereinstellungen einen TPU-Slice fest, reserviert Ray Serve für jedes Replikat eine topologiebewusste SlicePlacementGroup und startet darüber die vLLM-Engine-Worker. Das Beispiel im Beitrag stellt google/gemma-4-26B-A4B-it auf einem TPU-V6E-Slice mit 4x4-Topologie und einer Tensor-Parallel-Größe von 16 bereit.

Was das für Plattformteams bedeutet

Anyscale versteht die Veröffentlichung als Erweiterung des Open-Source-Stacks PARK: PyTorch, KI-Modelle, Ray und Kubernetes. Dem Unternehmen zufolge behalten PyTorch-Entwickler den Eager-Execution-Ansatz und die Ray-APIs, die sie bereits kennen, während Plattformteams von der topologiebewussten Autoskalierung, der Ausfallsicherheit und der Governance über mehrere Beschleuniger hinweg in GKE profitieren.

Teams auf GKE können nun verteiltes PyTorch-Training, vLLM-Inferenz und Reinforcement-Learning-Workloads über GPUs und TPUs hinweg auf einer einzigen Kubernetes-Steuerungsebene einplanen, skalieren und betreiben. Anyscale verweist auf durchgängige GKE-Tutorials zum Bereitstellen von LLMs und zum Trainieren von Modellen mit Ray und TorchTPU. Zudem will das Unternehmen die Grundlage auf Frameworks für RL-Post-Training in großem Maßstab sowie auf Slices der nächsten Generation von Ironwood (TPU v7x) auf GKE ausweiten.