Hugging Face kündigt native Unterstützung für TPUs, die von Google entwickelten Beschleunigerchips, in seinem Ökosystem an. Dank des Backends TorchTPU für PyTorch laufen Transformers, Diffusers und Kernels direkt auf diesen Prozessoren, wie das Unternehmen in einem Blogbeitrag mitteilte.
- Hugging Face unterstützt TPUs von Google nativ in seinem Ökosystem.
- Das Backend TorchTPU lässt Transformers, Diffusers und Kernels direkt auf TPUs laufen.
- Optimum TPU wird dafür nicht mehr benötigt.
- 96,5 % der Tests auf TPU bestehen.
Ein TPU, kurz für Tensor Processing Unit, ist ein Chip, der entwickelt wurde, um Matrixberechnungen zu beschleunigen, die den Großteil der Zeit von neuronalen Netzwerken in Anspruch nehmen. Google entwickelt diese Chips seit etwa einem Jahrzehnt und nutzt sie zum Trainieren und Betreiben seiner eigenen Modelle, darunter Gemini. Sie sind auf Google Cloud verfügbar, von einzelnen Chips bis hin zu Clustern mit mehreren Tausend.
Bisher hatte sich die Softwareumgebung für TPUs um TensorFlow und JAX aufgebaut, und PyTorch-Nutzer mussten größeren Aufwand betreiben. Um PyTorch auf TPU auszuführen, musste man PyTorch/XLA (torch_xla) verwenden, eine Lösung, die im Blogbeitrag als funktional aber umständlich beschrieben wird: Tensoren wurden lazy evaluiert, sodass der Code nicht immer zum erwarteten Zeitpunkt ausgeführt wurde, und kleine Änderungen konnten Neukompilierungen des Graphen auslösen. Hugging Face hatte Optimum TPU entwickelt, um Inferenz und Training seiner Modelle auf TPU zu ermöglichen. Mit TorchTPU ist dieses Tool nicht mehr nötig, um Modelle von Transformers und Diffusers laufen zu lassen.
TorchTPU, ein natives Backend für PyTorch
TorchTPU fügt die TPU als neues Backend zu PyTorch hinzu, ohne den üblichen PyTorch-Code zu ändern. Das Backend umfasst für TPU optimierte Pallas-Kernel, Unterstützung für verteiltes Training und eine Integration mit torch.compile für Nutzer, die nach noch mehr Geschwindigkeit suchen. Die Installation auf einer TPU-VM erfolgt mit einem einzigen Befehl: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. Der Beitrag präzisiert, dass dieser Index die CPU-Version von torch bereitstellt, die leichter ist als die CUDA-Version, die pip standardmäßig herunterlädt und die eine TPU nicht benötigt.
Transformers, Diffusers und Kernels auf TPU
In Transformers v5, dessen Ziel es ist, die Referenz für Modelldefinitionen zu sein, damit sie überall laufen, erfolgt das Laden eines Modells auf die TPU über den Parameter device_map= »tpu ». Der Beitrag veranschaulicht diesen Schritt mit dem Modell Qwen/Qwen3.5-9B. Standardmäßig läuft das Modell im Eager-Modus, kann aber auch torch.compile oder eine verteilte Konfiguration nutzen.
In Diffusers ist die TPU nun ein First-Class-Gerät. Der für Prototyping geeignete Eager-Modus benötigt nur einen Aufruf von pipe.to(« tpu »), während der für den Einsatz gedachte kompilierte Modus einen zusätzlichen Kompilierungsaufruf hinzufügt. Für Modelle, die zu groß für einen einzelnen Chip sind, verteilt eine an from_pretrained übergebene TensorParallelConfig den Transformer beim Laden auf mehrere Chips. Der Beitrag zeigt Beispiele, die auf einer Cloud TPU v6e-4 mit den Modellen MiniMax-H3 für Video und Audio sowie Qwen-Image-2.1 ausgeführt wurden.
Die Kernels-Bibliothek bietet Zugriff auf einsatzbereite TPU-Kernel, geschrieben in Pallas, das eine Low-Level-Kontrolle über die Hardware bietet. Auch Kernel aus der Tokamax-Bibliothek können integriert werden. Laut dem Beitrag ist bei einem Batch von 64 die Decodierung auf Qwen3-30B-A3B 21,5-mal und auf Gemma 4 26B-A4B 12,8-mal schneller als bei der Standardimplementierung, mit einer TPU v6e-8.
Kontinuierliche Tests
Hugging Face gibt an, die Testsuite von Transformers kontinuierlich auf TPU auszuführen und an der Einrichtung einer nächtlichen Continuous Integration auf diesen Chips zu arbeiten, deren Ergebnisse auf einem dedizierten Dashboard veröffentlicht würden. Laut dem Beitrag bestehen 96,5 % der auf TPU ausgeführten Tests. Jeder Fehlschlag führt zu einem Reproduktionsskript und einer Meldung im TorchTPU-Repository: 34 Probleme wurden gemeldet, davon sind 19 bereits behoben. Diese Tests haben auch Bugs in Transformers aufgedeckt, insbesondere in StaticCache und in den Tensor-Parallelismus-Plänen von Qwen3, die bereits im Voraus behoben wurden. Bei Problemen werden Nutzer aufgefordert, ein Issue zu öffnen.
