Hugging Face bringt native TPU-Unterstützung über das PyTorch-Backend TorchTPU

Von torch_xla zum nativen Backend

TPU steht für Tensor Processing Unit. Laut Hugging Face handelt es sich um Googles eigene KI-Beschleuniger, die für die Matrixrechnung gebaut sind, mit der neuronale Netze den Großteil ihrer Zeit verbringen. Google entwickelt sie seit rund zehn Jahren und nutzt sie, um eigene Modelle wie Gemini zu trainieren und bereitzustellen. Verfügbar sind sie in der Google Cloud, vom einzelnen Chip bis hin zu Pods mit Tausenden Chips.

Das Hindernis war nach Darstellung von Hugging Face immer die Software. Der TPU-Stack ist rund um TensorFlow und JAX gewachsen, weshalb PyTorch-Nutzer mehr Aufwand für den Einstieg hatten. Bislang bedeutete das PyTorch/XLA (torch_xla). Das funktionierte, doch Tensoren wurden verzögert (lazy) ausgewertet, sodass Code nicht immer dann lief, wenn man es erwartete. Kleine Änderungen konnten Neukompilierungen des Graphen auslösen, und das Debugging war schwer nachvollziehbar.

Hugging Face hatte Optimum TPU entwickelt, um das zu erleichtern, und bot damit Inferenz- und Trainingsunterstützung für seine Modelle auf TPUs. Mit der nativen Unterstützung durch TorchTPU ist Optimum TPU nach Angaben des Unternehmens nicht mehr nötig, um Hugging-Face-Modelle mit Transformers und Diffusers auszuführen.

TorchTPU fügt PyTorch die TPU als neues Backend hinzu, sodass Nutzer ganz normalen PyTorch-Code schreiben. Dazu kommen auf TPUs abgestimmte Pallas-Kernel, Unterstützung für verteiltes Training und eine torch.compile-Integration für mehr Tempo. Die Installation auf einer TPU-VM erfordert nur einen Befehl: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. Der zusätzliche Index liefert den CPU-Build von torch, da pip sonst den deutlich größeren CUDA-Build zieht, der auf einer TPU nicht gebraucht wird.

Transformers, Diffusers und Kernels auf der TPU

In Transformers genügt es, beim Laden eines Modells auf einer TPU device_map="tpu" zu übergeben. Das Beispiel von Hugging Face lädt auf diese Weise Qwen/Qwen3.5-9B und generiert Text. Das Modell läuft standardmäßig im Eager-Modus, optional sind torch.compile oder ein verteiltes Setup möglich.

In Diffusers ist die TPU nun ein vollwertiges Device. Für den Eager-Modus zum Prototyping oder Debugging reicht pipe.to("tpu"), für den kompilierten Modus im Produktivbetrieb kommt ein einziger Compile-Aufruf hinzu. Pipelines und Code bleiben laut Hugging Face unverändert, einen TPU-Fork muss niemand pflegen. Der erste Aufruf kompiliert, spätere Aufrufe mit derselben Bildgröße verwenden den Graphen erneut.

Für Modelle, die zu groß für einen Chip sind, verteilt eine an from_pretrained übergebene TensorParallelConfig den Transformer beim Laden auf mehrere Chips. Viele gängige Pipelines unterstützen das von Haus aus, bei den übrigen ist es „nur einen tp_plan entfernt“. Auf einer Cloud TPU v6e-4 mit TP=4, bf16 und Batchgröße 1 nennt der Beitrag als Beispiel ein Bild, das kompiliert in 3,5 s statt im Eager-Modus in 8,0 s erzeugt wird (1024×1024, 40 Schritte), sowie einen Video-und-Audio-Clip mit MiniMax-H3 in 188 s kompiliert gegenüber 262 s im Eager-Modus.

Der Kernels Hub bietet außerdem einsatzbereite, in Pallas geschriebene TPU-Kernel, und auch Kernel aus der Tokamax-Bibliothek lassen sich nutzen. Wird ein Tokamax-Ragged-Dot-Kernel in ein Mixture-of-Experts-Modell eingebunden, ist das Decoding bei Batchgröße 64 laut Hugging Face auf Qwen3-30B-A3B 21,5-mal und auf Gemma 4 26B-A4B 12,8-mal schneller. Die Werte vergleichen den Hub-Kernel mit der Standardimplementierung auf einer TPU v6e-8.

Testergebnisse und was das für Nutzer bedeutet

Hugging Face will die TPU-Unterstützung verbessern, indem die Testsuite von Transformers auf der TPU läuft, und arbeitet daran, eine nächtliche CI auf TPU mit Ergebnissen im Transformers-Dashboard einzurichten. Das Unternehmen bezeichnet das als laufende Arbeit: Von den Tests, die auf der TPU laufen, bestehen 96,5 %. Jeder Fehlschlag wird zu einem kleinen Reproduktionsskript und einem Issue im TorchTPU-Repository. Bisher wurden 34 Issues gemeldet, 19 davon sind bereits behoben. Der Betrieb auf der TPU hat zudem Fehler in Transformers selbst zutage gefördert, unter anderem in StaticCache und in den Qwen3-Tensor-Parallel-Plänen, die im Upstream behoben wurden.

Für PyTorch-Nutzer, die Zugang zu Google-Cloud-TPUs wollen, ändert sich praktisch Folgendes: Bestehende Hugging-Face-Codepfade lassen sich mit einem Device-String auf die Hardware ausrichten, statt einen separaten Software-Stack zu benötigen. Der Anspruch von Hugging Face lautet: „TPUs sind für PyTorch-Nutzer keine Nebenbaustelle mehr.“ Das Unternehmen verweist auf TorchTPU auf GitHub, auf TPU-Kernel im Hub und auf die Cloud-TPU-Dokumentation und bittet Nutzer, bei Problemen ein Issue zu eröffnen.