PyTorch auf TPU und Trainium: Wie weit sinkt die Nvidia-Abhängigkeit?

PyTorch auf TPU und Trainium braucht dank TorchTPU und TorchNeuron nur noch einen Gerätewechsel. Wir prüfen Zahlen, Grenzen, echte Kosten und wer jetzt testen sollte.

Bericht

PyTorch auf Googles TPUs oder AWS-Trainium-Chips zu betreiben, erforderte bisher eine Zwischenbibliothek und einigen Umbau. Am 19. und 20. Oktober 2026, kurz vor der PyTorch Conference, haben beide Anbieter ein natives PyTorch-Backend als Open Source veröffentlicht: TorchTPU (Google, gemeinsam mit Metas PyTorch-Team entwickelt) und TorchNeuron (AWS, öffentliche Beta). Das gemeinsame Ziel: PyTorch-Code für Nvidia-GPUs soll mit sehr wenigen Änderungen auf ihren Chips laufen.

• Gerät wechseln statt neu schreiben. Bei TorchNeuron ist laut AWS die einzige Änderung model.to('neuron'). Für TorchTPU verspricht Google „minimale Codeänderungen“; in Hugging Face Transformers v5 genügt device_map="tpu".
• Beide Ausführungsmodi. Beide unterstützen den Eager-Modus (Operation für Operation, gut zum Debuggen) und torch.compile (Kompilierung des ganzen Graphen für die Produktion).
• Verteiltes Training inklusive. DDP, FSDP2 und DTensor laufen auf beiden; TorchNeuron ergänzt Tensor- und Pipeline-Parallelität.
• Anbindung an das Serving-Ökosystem. TorchTPU bringt Backends für vLLM und SGLang mit; TorchNeuron bietet vLLM Neuron sowie Unterstützung in Transformers v5, Diffusers und trl.
• Open Source unter Apache 2.0 für beide Projekte auf GitHub.

Analyse

Analyse: Reicht eine Codezeile, um Nvidia-GPUs zu verlassen?

Kurz gesagt: Sie reicht zum Testen, nicht zum Umziehen. Die eigentliche Neuerung liegt in der Art der Integration; die Ersparnis hängt stark von Ihren Modellen ab.

Was ist anders als bisher?

Der entscheidende Unterschied ist PrivateUse1, die offizielle Erweiterungsschnittstelle von PyTorch für neue Hardware. Ein Tensor auf TPU oder Trainium ist ein normaler PyTorch-Tensor ohne Wrapper. Deshalb brauchen Bibliotheken wie Hugging Face, vLLM oder TorchTitan keinen eigenen Zweig pro Chip mehr.

Warum läuft das ohne Neuschreiben?

Beim Aufruf von torch.compile erfasst PyTorch mit TorchDynamo den Rechengraphen. Statt ihn an einen GPU-Compiler zu geben, übersetzt TorchTPU ihn nach StableHLO und lässt XLA für die TPU optimieren. TorchNeuron geht einen ähnlichen Weg über Torch-MLIR, StableHLO und den Neuron-Compiler. Entwickler schreiben weiter gewohntes PyTorch; die Übersetzung auf den Chip passiert darunter.

Der Modus „Fused Eager“ von TorchTPU fasst viele kleine Operationen automatisch zu größeren Blöcken zusammen, bevor sie zum Chip gehen. Google nannte Zuwächse von 50 % bis über 100 % gegenüber der Einzelausführung. Das ist eine Herstellerzahl ohne unabhängige Messung.

Wen betrifft es und wie?

• ML-Entwickler: können dieselbe Trainingsschleife auf GPU, TPU und Trainium vergleichen. Eigene Kernel erfordern weiterhin Pallas (TPU) oder NKI (Trainium).
• Technische Führungskräfte: gewinnen Verhandlungsmacht bei GPU-Preisen. Ein Cloud-Vertrag hängt nicht mehr an CUDA, wenn Modelle auf mehreren Chips laufen.
• Kleine und mittlere Unternehmen: nutzen Modelle meist per API, der Effekt ist also indirekt. Teams, die offene Modelle (Qwen, Llama) selbst auf Google Cloud oder AWS betreiben, profitieren direkt.
• Marketing und Produktteams: sofort ändert sich nichts. Ein Vorteil käme über günstigere API-Preise von Anbietern, die diese Chips nutzen.

Was kostet das in der Praxis?

Keiner der Anbieter hat neue Preise veröffentlicht. Die einzige Kostenzahl stammt von Deepgram: rund 60 % mehr verarbeitetes Audio pro Dollar als mit den bisherigen GPUs. Die Tabelle ist ein hypothetisches Beispiel, das diese Zahl auf zwei Größenordnungen anwendet, mit angenommenen Migrationskosten von 30.000 USD (etwa zwei Ingenieure für 1–2 Monate).

Fazit: Bei kleiner Rechnung lohnt sich ein Chipwechsel selten. Bei großer Rechnung und stabilem Modell beginnt sich die Rechnung zu tragen. Die 60 % stammen von einem Kunden und einer Last (Audio); Ihr Modell kann besser, schlechter oder gar nicht profitieren.

Risiken und Grenzen

• Unvollständige Operator-Abdeckung. Im PyTorch-Testlauf vom August 2026 wurden rund 22.000 Tests übersprungen, weil TorchNeuron der passende Operator fehlte. Die erste Version unterstützt keine dynamischen Shapes, keine Sparse Tensors und kein NestedTensor; FP8 ist in privater Beta.
• Dynamische Shapes bleiben eine TPU-Schwäche. Google räumt ein, dass wechselnde Sequenzlängen oder Batchgrößen Neukompilierungen auslösen und die erste Iteration hohe Latenz hat. Ein Chatbot mit stark schwankenden Prompt-Längen ist genau die betroffene Last.
• Die Abhängigkeit verschiebt sich nur. TPUs gibt es nur in Google Cloud, Trainium nur bei AWS. Weg von CUDA kann enger an eine Cloud binden.
• Reiner CUDA-Code läuft nicht. CUDA Graphs, NCCL-spezifische Aufrufe und eigene CUDA-Kernel müssen neu geschrieben werden.

Was noch zu prüfen ist

• Die 52 % MFU und 8.192 Chips stammen aus internen TorchTitan-Tests von Google, ohne unabhängigen Vergleich mit preisgleichen GPUs.
• Die Zahlen von Zenteiq (etwa Faktor 4) und Deepgram (+60 % pro Dollar) sind von den Anbietern ausgewählte Kundenzitate.
• Die von TechCrunch zitierte Angabe „85 % der KI-Apps nutzen PyTorch“ hat keine benannte Quelle; TBN Grow konnte sie nicht verifizieren.
• So prüfen Sie selbst: Nehmen Sie ein produktives Modell, messen Sie Tokens/Sekunde und Kosten pro 1 Mio. Tokens auf Ihren GPUs und betreiben Sie dasselbe Modell auf TPU v6e oder Trainium2 mit gleichen Eingabelängen.

Einschätzung von TBN Grow

Es handelt sich um eine Änderung der Infrastruktur-Software mit echtem Nutzen für Teams, die offene Modelle auf Google Cloud oder AWS trainieren oder selbst betreiben. Wer monatlich GPU-Rechnungen im fünfstelligen Dollarbereich hat, sollte jetzt mit einem stabilen Modell testen. Kleine Teams, API-Nutzer oder Lasten mit stark schwankenden Eingaben sollten ein bis zwei Versionen abwarten, bis Operator-Abdeckung und dynamische Shapes ausgereift sind.

Ausblick

Worauf Sie bei PyTorch auf TPU und Trainium achten sollten

Was Sie sofort tun können

• [ ] Suchen Sie jeden direkten cuda-Aufruf im Code und ersetzen Sie ihn durch die torch.accelerator-API, um Chipwechsel zu erleichtern.
• [ ] TorchTPU testen: pip install torch_tpu==0.2.0 auf einer TPU-v6e-VM (laut README der einfachste Einstieg; mit --extra-index-url https://download.pytorch.org/whl/cpu, damit nicht die CUDA-Version installiert wird).
• [ ] TorchNeuron testen: pip install torch-neuronx auf einer trn2-Instanz und Gerät auf neuron umstellen.
• [ ] Die Tests Ihres eigenen Projekts laufen lassen, um nicht unterstützte Operatoren zu finden, bevor Sie Kosten schätzen.
• [ ] Kosten pro 1 Mio. Tokens oder pro Audiostunde vergleichen, nicht nur die Geschwindigkeit.

Zeitplan und noch nicht Angekündigtes

• AWS will TorchAO und FP8 in der nächsten Version allgemein verfügbar machen und Operator-Lücken schrittweise schließen; ohne Termine.
• Google plant dynamische Shapes in torch.compile und eine Bibliothek vorkompilierter Kernel; ohne Zeitplan.
• Google hat TorchTPU weder als Beta noch als GA gekennzeichnet und keine vollständige Liste unterstützter TPU-Generationen veröffentlicht.
• Nvidia hat sich offiziell nicht geäußert.

Weiterlesen: Chinas KI-Chipmarkt im Wandel: Nvidia bei 34 %, Huawei bei 28 % und die Renditefrage bei Investitionen in KI-Infrastruktur.

Prüfquellen: Google Cloud Blog, Google Developers Blog, TorchTPU auf GitHub, PyTorch Blog – TorchNeuron, AWS What's New, TechCrunch.