Models

TorchTPU: PyTorch, Transformers und Diffusers laufen jetzt nativ auf TPUs

models pytorch tpu huggingface transformers

TPUs ohne Umweg

Googles Tensor Processing Units waren für PyTorch-Nutzer bisher eine Hürde: Der Weg führte über PyTorch/XLA mit lazy Tensors, unerwarteten Neukompilierungen und schwer nachvollziehbarem Debugging. Hugging Face hatte dafür Optimum TPU gebaut. Mit TorchTPU wird die TPU nun ein reguläres PyTorch-Backend – Optimum TPU wird damit nicht mehr benötigt.

Was neu ist

  • Installation auf einer TPU-VM per einzelnem pip-Befehl (mit CPU-Index für torch, damit nicht der große CUDA-Build gezogen wird)
  • Transformers: Modelle laden mit device_map="tpu", standardmäßig im Eager-Modus, optional mit torch.compile oder verteilt
  • Diffusers: pipe.to("tpu") für Prototypen, ein Compile-Aufruf für den Produktivbetrieb. Große Modelle lassen sich per TensorParallelConfig beim Laden über mehrere Chips shardieren
  • Kernels Hub: fertige Pallas-Kernels, auch aus der Tokamax-Bibliothek. Ein Ragged-Dot-Kernel beschleunigt das Decoding bei Batch-Größe 64 um das 21,5-Fache bei Qwen3-30B-A3B und 12,8-Fache bei Gemma 4 26B-A4B (TPU v6e-8)

Reifegrad

Hugging Face lässt die Transformers-Testsuite auf TPU laufen: 96,5 % der ausgeführten Tests bestehen. 34 Probleme wurden an das TorchTPU-Repo gemeldet, 19 sind behoben; zudem wurden Bugs in Transformers selbst gefunden (StaticCache, Qwen3-Tensor-Parallel-Pläne). Ein nächtliches TPU-CI ist in Arbeit.

Fazit: Wer PyTorch-Code hat und Cloud-TPUs nutzen will, muss ihn kaum noch anpassen. Grün ist noch nicht alles – aber der Einstieg ist deutlich einfacher geworden.

Quelle: Hugging Face Blog