Google macht TorchTPU öffentlich: PyTorch läuft nativ auf TPUs

Worum es geht

Google hat TorchTPU als Open-Source-Projekt auf GitHub veröffentlicht (google-pytorch/torch_tpu, Apache-2.0-Lizenz). Es ist ein natives PyTorch-Backend für Googles Tensor Processing Units (TPUs), also die eigenen KI-Beschleunigerchips des Konzerns. Parallel hat Hugging Face einen Blogbeitrag veröffentlicht, wonach die Bibliotheken Transformers, Diffusers und Kernels nun direkt auf TPUs laufen. Wer ein Modell aus dem Hub nutzen will, schickt es an das Gerät "tpu".

Das Problem bisher

TPUs sind Chips, die auf die Matrixrechnungen zugeschnitten sind, die neuronale Netze größtenteils ausmachen. Google entwickelt sie seit etwa zehn Jahren und nutzt sie nach Hugging-Face-Angaben für Training und Betrieb eigener Modelle, darunter Gemini. Verfügbar sind sie in der Google Cloud, vom Einzelchip bis zu Pods mit tausenden Chips.

Das Hindernis war die Software: Der TPU-Stack wuchs um TensorFlow und JAX herum. PyTorch-Nutzer mussten über PyTorch/XLA (torch_xla) gehen. Das funktionierte, hatte aber Eigenheiten: Tensoren wurden „lazy“ ausgewertet, Code lief also nicht unbedingt dann, wann man es erwartete. Kleine Änderungen konnten das Neukompilieren des Rechengraphen auslösen, und das Debuggen war schwer nachvollziehbar. Hugging Face hatte deshalb Optimum TPU gebaut; mit TorchTPU ist das laut dem Blog für Transformers und Diffusers nicht mehr nötig.

Was TorchTPU bietet

Laut README und Blog enthält das Projekt:

• eigene ATen-Kernel (die Grundoperationen von PyTorch) für TPUs
• einen Compilation Cache, um XLA-Berechnungen schneller auszuführen
• verteiltes Training mit Daten- und Tensorparallelität (Tensorparallelität teilt ein Modell auf mehrere Chips auf)
• Pallas-Kernel, optimiert für TPUs
• Integration mit torch.compile für den Graph-Modus
• Beispiele für Llama, Qwen, ResNet, DLRM, minGPT und UNet sowie Werkzeuge für Debugging und Benchmarking

Eine Anleitung gibt es auf torchtpu.readthedocs.io; die interaktiven Notebooks sind Marimo-Notebooks (ähnlich Jupyter) und sollen auf einer TPU-VM laufen, am einfachsten auf einer v6e-GCE-Instanz. Der Code ist selbst für Python 3.11+ ausgelegt.

Installation und Nutzung

Auf einer TPU-VM genügt:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

Der Zusatzindex liefert die CPU-Variante von torch; ohne ihn lädt pip die deutlich größere CUDA-Version, die auf einer TPU nicht gebraucht wird. Alternativ lässt sich das Wheel mit Bazel aus dem Quellcode bauen. Das README nennt dafür Details wie die Option --config=wheel_common und --config=no_rbe, falls man keine Zugangsdaten für Googles Remote Build Execution hat.

Integration in Hugging Face

In Transformers v5 genügt device_map="tpu" in from_pretrained. Das Beispiel im Blog lädt Qwen/Qwen3.5-9B direkt auf die TPU. Standardmäßig läuft das Modell im Eager-Modus (Operation für Operation); mit torch.compile ist es schneller, auch verteilte Setups sind möglich.

In Diffusers ist TPU laut Blog ein vollwertiges Gerät: Für Prototypen reicht pipe.to("tpu"), für den Produktivbetrieb kommt ein compile-Aufruf dazu. Der erste Aufruf kompiliert, weitere Aufrufe mit gleicher Bildgröße nutzen den Graphen wieder. Für zu große Modelle übergibt man eine TensorParallelConfig, und das Modell wird beim Laden auf mehrere Chips verteilt.

Beispiele auf einer Cloud TPU v6e-4 mit TP=4, bf16, Batch 1:

Für Qwen-Image-2.1 wurde zudem ein 2048×2048-Bild gezeigt, ohne Zeitangabe.

Kernels und Tests

Der Kernels Hub bietet fertige TPU-Kernel, geschrieben in Pallas. Auch Kernel aus der Tokamax-Bibliothek lassen sich einbinden. Im Beispiel wird ein Ragged-Dot-Kernel in ein Mixture-of-Experts-Modell (Qwen3-30B-A3B-Instruct-2507) eingesteckt. Bei Batchgröße 64 ist das Decoding laut Hugging Face 21,5-mal schneller bei Qwen3-30B-A3B und 12,8-mal schneller bei Gemma 4 26B-A4B (TPU v6e-8, Vergleich mit der Standardimplementierung). Je größer der Batch, desto größer der Gewinn.

Hugging Face lässt die Transformers-Testsuite auf TPU laufen und baut eine nächtliche CI auf. Von den Tests, die auf TPU laufen, bestehen 96,5 %. Fehler werden als kleine Reproduktionsskripte und Issues im TorchTPU-Repo gemeldet: 34 bisher, 19 behoben. Zudem fand man Bugs in Transformers selbst (StaticCache, Tensor-Parallel-Pläne von Qwen3) und korrigierte sie.

Einordnung

TorchTPU senkt die Hürde für PyTorch-Nutzer erheblich, TPUs zu nutzen, ohne Code nach JAX zu portieren. Das ist vor allem für Teams interessant, die in der Google Cloud Rechenleistung für Training oder Serving brauchen und bei NVIDIA-Hardware knappe Kapazitäten haben. Die Geschwindigkeitswerte stammen von Hugging Face selbst und vergleichen mit der Standardimplementierung, nicht mit GPUs; Vergleiche zu Kosten oder zu anderen Beschleunigern fehlen. Die Unterstützung ist noch nicht vollständig: Nicht alle Tests laufen, die Abhängigkeitsliste im README ist noch „TBD“, und die Version ist 0.2.0. Zum Ausprobieren ist eine TPU-VM in der Google Cloud nötig, was Kosten verursacht.