Wenn Sie schon einmal versucht haben, PyTorch auf einer TPU auszuführen, haben Sie Ihren Code wahrscheinlich portiert auf torch_xla, oder Sie mussten am Ende JAX lernen. Mit TorchTPU, ein natives PyTorch-Backend, ist TPU jetzt einfach nur ein weiteres Gerät. Dies eröffnet native TPU-Unterstützung im Hugging Face-Ökosystem. Transformers, Diffusers und Kernels laufen nativ auf TPU. Wählen Sie ein Modell vom Hub aus und senden Sie es an "tpu", und fertig.
Moment, was ist ein TPU? 🤔
TPU steht für Tensor Processing Unit. TPUs sind Googles eigene KI-Beschleuniger. Es sind Chips, die für eine einzige Aufgabe gebaut wurden: die Matrixmathematik, auf die neuronale Netze den Großteil ihrer Zeit verwenden. Google entwirft sie seit etwa zehn Jahren und nutzt sie, um seine eigenen Modelle zu trainieren und bereitzustellen, einschließlich Gemini.
TPUs sind verfügbar auf Google Cloud, von einem einzelnen Chip bis hin zu Pods mit Tausenden von Chips. Sie glänzen, wenn Sie viel Rechenleistung benötigen, sei es für Training oder für Serving im großen Maßstab.
Der Haken war immer die Software. Der TPU-Stack entstand rund um TensorFlow und JAX, und PyTorch-Nutzer mussten sich mehr anstrengen, um Zugriff zu bekommen.
Der alte Weg
Bisher bedeutete das Ausführen von PyTorch auf TPUs den Weg über PyTorch/XLA (torch_xla). Es funktionierte, aber es hatte seine Eigenheiten. Tensoren waren lazy (die Tensoren, nicht Sie), daher lief Ihr Code nicht immer dann, wenn Sie dachten, dass er es tut. Kleine Änderungen konnten Graphen-Neukompilierungen auslösen, und das Debuggen war schwer nachzuvollziehen.
Um das Leben leichter zu machen, haben wir Optimum TPU entwickelt. Es bot Hugging Face-Modellen Inferenz- und Trainingsunterstützung auf TPUs. Jetzt, da TorchTPU native Unterstützung bietet, ist Optimum TPU nicht mehr erforderlich, um Hugging Face-Modelle mit Transformers und Diffusers auszuführen.
TorchTPU tritt auf 🚀
TorchTPU fügt TPU als neues Backend zu PyTorch hinzu. Sie schreiben ganz normalen PyTorch-Code. Es wird mit für TPUs optimierten Pallas-Kernels geliefert, unterstützt verteiltes Training und bietet eine torch.compile -Integration für den Fall, dass Sie mehr Geschwindigkeit wünschen. 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 Ihnen den CPU-Build von torch. Ohne ihn lädt pip den viel größeren CUDA-Build, den Sie auf einer TPU nicht benötigen.
Ihre Lieblingsbibliotheken, jetzt auf TPU
Transformers
Transformers v5 wurde um eine einzige Idee herum gebaut: die maßgebliche Quelle für Modelldefinitionen zu sein, damit sie überall laufen können. „Überall“ umfasst jetzt auch TPUs. Um ein Modell auf einer TPU zu laden, übergeben Sie device_map="tpu". Das ist alles.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu", # will load model directly on the TPU
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))
Das Modell läuft standardmäßig im Eager-Modus. Sie können auch torch.compile verwenden, um eine bessere Leistung zu erzielen, oder es in einem verteilten Setup ausführen.
Bilder? Videos? Diffusers malt sie auf TPUs 🎨
TPU ist jetzt ein erstklassiges Gerät in Diffusers. Für Prototyping oder Debugging benötigt der Eager-Modus nur pipe.to("tpu"). Für den Einsatz fügt der kompilierte Modus einen einzigen compile -Aufruf hinzu. Ihre Pipelines und Ihr Code bleiben gleich, und es gibt keinen TPU-Fork zu pflegen.
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu") # eager mode: that's the whole change
# optional: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]
Der erste Aufruf kompiliert, und jeder weitere Aufruf mit derselben Bildgröße verwendet den Graphen wieder.
Modelle zu groß für einen Chip? Übergeben Sie einen TensorParallelConfig an from_pretrained , und der Transformer wird beim Laden über die Chips geshardet, sodass das vollständige Modell nie an einem einzigen Ort passen muss. Viele beliebte Pipelines unterstützen dies sofort, und die übrigen sind nur einen tp_plan entfernt. Der Diffusers-TPU-Leitfaden enthält das vollständige Rezept.
Hier sind einige Beispiele, die auf einer Cloud TPU v6e-4 ausgeführt wurden:
TP=4 (jedes Modell auf vier Chips aufgeteilt) · bf16 · Batch 1
| MiniMax-H3 (Video + Audio) | Qwen-Image-2.1 | |
|---|---|---|
| "Das flauschige Plüsch-Maskottchen von Hugging Face aus dem Bild, ein rundes gelbes lächelndes Gesicht mit rosigen Wangen und zwei kleinen offenen Händen, steht am Bug einer antiken griechischen Triere […] Klang von brandenden Wellen, knarrendem Holz, Wind in der Takelage, fernen Möwen und einer leisen epischen Orchesterwelle." · 1344×768, 5.2 s, 49 Schritte · 188 s kompiliert, 262 s eager | "Ein Capybara mit Zaubererhut, das bei Kerzenlicht ein Buch liest, Ölgemälde" · 1024×1024, 40 Schritte · 3.5 s kompiliert, 8.0 s eager | "Ein Leuchtreklameschild mit der Aufschrift 'QWEN IMAGE 2.1', regnerische Nacht, Reflexionen auf nassem Pflaster" · 2048×2048, 40 Schritte |
Kernels
Der Kernels Hub bietet einsatzbereite TPU-Kernels. Laden Sie einen mit der Kernels-Bibliothek und geben Sie Ihrem Modell einen Geschwindigkeitsschub.
TPU-Kernel werden in Pallasgeschrieben, was Ihnen Kontrolle auf niedriger Ebene über die Hardware gibt. Sie können auch Kernel aus der Tokamax -Bibliothek einbinden. Hier ist ein Ragged-Dot-Kernel aus Tokamax, eingebaut in ein Mixture-of-Experts-Modell:
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-30B-A3B-Instruct-2507",
experts_implementation="tokamax_ragged_dot",
)
Je größer der Batch, desto größer der Gewinn. Bei einer Batch-Größe von 64 ist die Dekodierung auf Qwen3-30B-A3B 21.5× schneller und auf Gemma 4 26B-A4B 12.8× schneller.
Decode-Durchsatz auf einer TPU v6e-8, Hub-Kernel gegenüber der Standard-Implementierung.
Wir testen es, damit Sie es nicht müssen 🧪
Wir verbessern die TPU-Unterstützung in transformers kontinuierlich, indem wir die Test-Suite der Bibliothek ausführen. Wir arbeiten daran, eine nächtliche CI auf TPU einzurichten und die Ergebnisse im (transformers dashboard)[ verfügbar zu machen.https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].
Das Transformers-CI-Dashboard, mit TPU-Ergebnissen neben anderen Beschleunigern. In Arbeit!
Noch ist nicht alles grün, und das ist uns recht. Von den Tests, die auf TPU laufen, bestehen 96.5%. Jeder Fehler, den wir finden, wird zu einem kleinen Repro-Skript und einem Issue im TorchTPU-Repo. Bisher haben wir 34 Issues gemeldet, und 19 sind bereits behoben. Der Betrieb auf TPU hat auch einige Bugs in Transformers selbst aufgedeckt, etwa in StaticCache und den Qwen3-Tensor-Parallel-Plänen. Wir haben diese Upstream behoben.
Wenn auf TPU etwas bricht, wollen wir es herausfinden, bevor Sie es tun.
Loslegen
TPUs sind für PyTorch-Nutzer nicht mehr nur ein Nebenschauplatz. Installieren Sie TorchTPU, holen Sie sich ein Modell vom Hub und senden Sie es an "tpu".
Wir können es kaum erwarten zu sehen, was Sie bauen 🤗. Falls Sie auf ein Problem stoßen, öffnen Sie ein Issue und informieren Sie uns.