Google und Meta starten TorchTPU, ein natives PyTorch-Backend für TPUs

Ein natives Backend statt JAX-Umweg

TPUs waren lange an die Ökosysteme von TensorFlow und JAX gebunden. TorchTPU, ein Projekt, das Mitte 2025 bei Google gemeinsam mit dem PyTorch-Team von Meta gestartet wurde, soll das ändern: Entwickler sollen PyTorch direkt auf der Hardware ausführen können.

Die Autoren beschreiben es als ein konsequent durchdachtes TPU-Backend für PyTorch. Sie grenzen es von früheren Ansätzen wie PyTorch/XLA und TorchAX ab, die ihrer Darstellung nach „PyTorch in die Arbeitsweise von JAX zwängen“. PyTorch/XLA, 2020 von Google eingeführt und auf Lazy Tensors basierend, erforderte laut dem Beitrag tiefgreifende Eingriffe in den Quellcode, um kompilierbare Bereiche zu definieren, litt unter erheblichen Latenzen bei der Kompilierung, unterstützte keine Eager-Ausführung und hinkte bei der Unterstützung neuer verteilter Primitive hinterher.

Bei TorchTPU hat die Eager-Ausführung höchste Priorität. Ein Modell zu portieren soll so einfach sein wie der Wechsel von .to("cuda") zu .to("tpu"). Die Autoren räumen ein, dass diese Änderung allein noch keine Spitzenleistung auf TPUs bringt, weil Shapes, Sharding und Operationen eines Modells auf eine GPU-Architektur zugeschnitten sein können, das Modell solle aber lauffähig sein.

Designprinzipien und aktueller Stand

Das Team hat sich drei Gruppen von Prinzipien gesetzt: PyTorch-Bürgerschaft („PyTorch citizenship“), Fokus auf den Nutzer und die Weiterentwicklung des Stands der Technik. In der Praxis bedeutet das: keine JAX-Abhängigkeiten in der Laufzeitumgebung, ein solider Eager-Modus sowie torch.compile auf Basis von Torch Dynamo statt Lazy Tensor, mit einer Leistung ähnlich wie bei JAX compile. Verteilte Idiome wie DDP, FSDP und DTensor werden unterstützt, indem TPU-Kollektive auf der untersten Ebene des PyTorch-Stacks bereitgestellt werden.

Laut dem Beitrag befindet sich TorchTPU in aktiver Entwicklung und wird von einigen der größten Google-Cloud-Kunden produktiv für Inferenz und Training eingesetzt. Es ist als Open-Source-Projekt auf GitHub verfügbar. Zum Zeitpunkt der Veröffentlichung nennt das Team folgenden Stand:

• Unterstützung für den Eager- und den kompilierten Modus, einschließlich eines neuen Eager-Modus namens DeferAndFuse.
• Unterstützung für rund 450 Aten-Operationen, darunter Forward- und Backward-Pass für alle Core-ATen-Operationen, die für die wichtigen Deep-Learning-Architekturen nötig sind.
• Entwicklung eigener Kernel über JAX/Pallas und Helion DSL.
• Validierte Ausführung ohne Anpassungen („out of the box“) für die 5.000 meistgeladenen Hugging-Face-Modelle, darunter LLMs wie Qwen, DeepSeek, Gemma, LLama und Mistral, Sprachmodelle wie Whisper, Bildmodelle wie CLIP und ViT sowie Text-Encoder wie BERT und T5.
• Validiertes DTensor-basiertes Sharding, Multi-Host-Inferenz mit vLLM und SGLang sowie Training mit TorchTitan unter Ray-basierter Orchestrierung.

Wie die TPU-Hardware das Design prägt

Der Beitrag enthält außerdem eine Einführung in die Hardware. TPUs sind spezialisierte ASICs mit TensorCores, die um Matrixmultiplikationseinheiten herum aufgebaut sind; einige Generationen ergänzen SparseCores für Sparse-Operationen und Kollektive. Der Beitrag nennt 288 GB High Bandwidth Memory mit 8,6 TB/s bei der TPU 8i sowie 19,2 TBps Inter-Chip-Interconnect bei TPU 8t und 8i.

Diese Eigenschaften erklären mehrere Designentscheidungen. Ein TensorCore verhält sich wie ein Single-Thread-Prozessor mit gewaltigen Matrixeinheiten, daher brauchen Modelle ausreichend breite Shapes, um die Spitzenleistung zu erreichen. Auf eine TPU kann nur ein Userspace-Prozess zugreifen, weshalb TorchTPU ein Modell mit einem Prozess pro Gerät verwendet, anders als das Modell von JAX mit einem Prozess pro Host. Folglich lässt sich RL-Colocation nicht ohne Weiteres unterstützen; das Team prüft nach eigenen Angaben Workarounds.

Der XLA-Compiler und die PjRT-Laufzeit setzen voraus, dass Shapes zur Kompilierzeit bekannt sind, was für PyTorch-Workloads eine Einschränkung darstellt. TorchTPU ergänzt deshalb begrenzte Dynamik („bounded dynamism“), bei der zur Kompilierzeit nur die Obergrenzen einer Dimension bekannt sein müssen. Kollektive Operationen werden als StableHLO-Operationen nach XLA übersetzt, statt eine Bibliothek wie NCCL zu verwenden. Da StableHLO keine In-Place-Mutationen erlaubt, stellt TorchTPU die gewohnte PyTorch-View-Semantik, einschließlich Mutationen, durch View-Dekomposition und -Inversion bereit.

Was das für PyTorch-Nutzer bedeutet

Für Teams, die bereits mit PyTorch arbeiten, bietet TorchTPU einen Weg, Google-TPUs auszuprobieren, ohne Code für JAX umzuschreiben. Die Autoren nennen als Ziel eine Lösung, die bei großen Workloads mit PyTorch auf GPUs konkurrieren kann, und bezeichnen die Arbeit als fortlaufend. Der hier vorliegende Quelltext bricht beim Autoload-Mechanismus ab und enthält nicht die von den Autoren erwähnten Benchmark-Ergebnisse. Leistungsaussagen bleiben daher vorerst Angaben des Projekts, bis diese Zahlen geprüft sind.