Google und Meta veröffentlichen TorchTPU, um PyTorch nativ auf TPUs auszuführen
Das Open-Source-Backend, das Mitte 2025 bei Google gemeinsam mit dem PyTorch-Team von Meta gestartet wurde, wird laut seinem Team bereits in der Produktion von einigen der größten Google-Cloud-Kunden eingesetzt.
KI-IllustrationGoogle und Meta haben TorchTPU vorgestellt, ein Open-Source-Backend, mit dem PyTorch nativ auf Googles TPU-Chips laufen kann. Das Team des Projekts erklärte, es sei Mitte 2025 bei Google in Zusammenarbeit mit dem PyTorch-Team von Meta entstanden und befinde sich in aktiver Entwicklung.
Laut dem Team wird TorchTPU bereits von einigen der größten Google-Cloud-Kunden in der Produktion für Inferenz- und Trainings-Workloads eingesetzt. TPUs seien historisch an die Ökosysteme von TensorFlow und JAX gebunden gewesen, so das Team.
Das Team beschrieb TorchTPU als Bruch mit früheren Ansätzen. Google veröffentlichte 2020 PyTorch/XLA, eine erste kompilierte PyTorch-Version für TPUs auf Basis von Lazy Tensors. Laut dem Team erforderte diese Integration tiefgreifende Änderungen am Quellcode, litt unter erheblicher Latenz bei der Kompilierung, unterstützte den Eager-Execution-Modus von PyTorch nicht und hinkte bei neueren Funktionen für verteiltes Rechnen hinterher. TorchTPU hingegen behandelt die Eager Execution als oberste Priorität und hat in seiner Laufzeit keine JAX-Abhängigkeiten.
Das erklärte Ziel ist, dass für die Umstellung einer Workload auf eine TPU im Code kaum mehr nötig sein soll, als das Zielgerät von einer GPU auf eine TPU zu ändern. Das Team warnte, dass eine solche Änderung allein nicht die maximale TPU-Leistung liefere, da Shapes, Sharding und Operationen eines Modells möglicherweise auf GPUs abgestimmt seien, das Modell solle aber dennoch sofort lauffähig sein.
Bisher hat das Projekt laut dem Team:
- Unterstützung für Eager- und kompilierte Modi hinzugefügt, darunter ein neuer Eager-Modus namens DeferAndFuse;
- Unterstützung für rund 450 ATen-Operationen hinzugefügt;
- benutzerdefinierte Kernel ermöglicht, die über JAX/Pallas und die Helion DSL geschrieben werden;
- die sofort einsatzbereite Ausführung der 5,000 meistheruntergeladenen Hugging-Face-Modelle validiert;
- die Integration mit vLLM, SGLang, TorchTitan, Ray und TorchRec validiert.
Das Team nannte auch Einschränkungen. Auf ein TPU-Gerät kann nur ein einziger User-Space-Prozess zugreifen, daher lässt sich Colocation für Reinforcement Learning nicht sofort unterstützen; das Team prüft nach eigenen Angaben Wege, dies zu umgehen. Außerdem wird Unterstützung für begrenzte Dynamik ergänzt, da der zugrunde liegende Compiler und die Laufzeit davon ausgehen, dass Tensor-Shapes zur Kompilierzeit bekannt sind.
Bei der Präzision verwendet TorchTPU standardmäßig einen mittleren Modus, der in bfloat16 multipliziert und in FP32 akkumuliert. Höhere Modi stellen mehr FP32-Präzision wieder her, bei etwa dem 3- bzw. 6-Fachen der Rechenkosten; das Team empfiehlt den Standard für die meisten Deep-Learning-Workloads.