Google hat TorchTPU als Open Source veröffentlicht. Die Software-Schicht ermöglicht es Entwicklern, PyTorch – ein weit verbreitetes Toolkit zum Erstellen und Ausführen von KI-Modellen – auf Googles Tensor Processing Units (TPUs) auszuführen. TorchTPU wurde gemeinsam mit Metas PyTorch-Team entwickelt und soll bestehende Workloads mit vergleichsweise wenigen Codeänderungen auf TPU-Hardware übertragen, statt sie neu schreiben zu müssen.
Laut Google unterstützt TorchTPU sowohl die eager Ausführung, bei der Operationen direkt beim Aufruf ausgeführt werden und die sich gut zum Debuggen eignet, als auch die kompilierte Ausführung für Produktions-Workloads. Entwickler können eine Geräteeinstellung von „cuda“ auf „tpu“ ändern und das TPU-Backend für verteilte Ausführung nutzen. Als Alternative verweist Google auf PyTorchs API zur Erkennung von Beschleunigern. TorchTPU ist auf GitHub verfügbar. Auf der Roadmap stehen jedoch noch dynamische Eingabeformen, kürzere Kompilierungszeiten, eine breitere Operator-Abdeckung sowie Reinforcement Learning und Quantisierung.
Was ändert sich für Entwickler?
Mit der Veröffentlichung werden TPUs für Software zugänglich, die viele PyTorch-Entwickler bereits nutzen. Laut Google funktioniert TorchTPU mit den Serving-Tools vLLM und SGLang sowie mit Trainings- und Orchestrierungsprojekten wie TorchTitan, Ray und TorchRec. Für Transformers v5 gibt Hugging Face an, dass Entwickler mit der Einstellung device_map="tpu" eine TPU auswählen können. Auch die Bibliothek Diffusers unterstützt die Ausführung auf TPUs, einschließlich des kompilierten Modus und der Aufteilung von Aufgaben auf mehrere Beschleuniger.
Damit verfolgt Google einen anderen Ansatz als mit PyTorch/XLA, der bisherigen Lösung des Unternehmens, um PyTorch auf TPUs auszuführen. Laut Google waren damit größere Codeänderungen nötig, die Kompilierungszeiten lang und der eager Modus von PyTorch nicht verfügbar. Google hat eine Testpipeline aufgebaut, die laut InfoQ rund 5.000 Hugging-Face-Modelle aus mehr als 500 Architekturen prüft. Das deutet auf umfangreiche Tests hin, bedeutet aber nicht, dass jedes Modell und jeder Workload unverändert funktioniert.
Wie schneidet TorchTPU im Vergleich zu anderer Hardware ab?
Google zufolge kann der Modus DeferAndFuse von TorchTPU – eine Mischung aus eager Ausführung und Kompilierung – mindestens 50% schneller sein als der Modus Strict Eager und zugleich den maximalen Speicherbedarf des Beschleunigers senken. In einem separaten Test führte Google ein Muse-Glimmer-Modell mit 30 Milliarden Parametern in Standard-PyTorch aus: InfoQ zufolge lag der Inferenzdurchsatz bei einer Sequenzlänge von 5,120 auf einem TPUv7-Chip bei 25,700 Tokens pro Sekunde und bei einer Länge von 10,240 bei 22,400. Beim Training auf vier TPUv7-Chips wurden bei diesen Sequenzlängen jeweils 5,600 und 4,600 Tokens pro Sekunde erreicht. Das sind von Google gemeldete Ergebnisse, kein unabhängiger Vergleich.
TechCrunch stellt TorchTPU neben Amazons TorchNeuron, das in derselben Woche für PyTorch-Workloads auf Amazons Trainium-Chips angekündigt wurde. Beide Projekte sollen die Nutzung von Alternativen zu Nvidia-GPUs erleichtern. Mit TorchTPU bekommen PyTorch-Nutzer eine weitere Hardware-Option. Die Veröffentlichung zeigt jedoch nicht, dass sich ein Workload dadurch günstiger ausführen lässt oder die bisherige Einrichtung eines Entwicklers übertrifft.
Warum es wichtig ist
Für Entwickler, die bereits PyTorch nutzen, könnte TorchTPU Tests oder den Einsatz von Workloads auf Google-TPUs erleichtern, ohne sie von Grund auf neu erstellen zu müssen. Ob sich der Wechsel bei einem bestimmten Modell in geringeren Kosten oder höherer Geschwindigkeit niederschlägt, ist noch offen; die genannten Leistungswerte stammen von Google selbst.
