TorchTPU-PyTorch-TPU-Leitfaden: Installation und Migration 2026

Ein cuda-zu-tpu-Port in PyTorch ist mit TorchTPU vergleichsweise überschaubar. Ob der Code anschließend zuverlässig und schnell läuft, entscheiden allerdings vor allem der TPU-Zugriff, das passende Paket, die Kompilierung und die verwendeten Tensorformen.

Was TorchTPU tatsächlich verspricht

TorchTPU ist ein natives PyTorch-Backend für Googles TPUs – keine gehostete Modell-API und auch kein Ersatz für eine TPU-VM. Die PyTorch-Ankündigung nennt Eager Execution, torch.compile, verteilte Primitive und die Unterstützung von rund 450 ATen-Operationen. Im öffentlichen Repository finden sich außerdem ATen-Kernels, Caching für Kompilierungen, Beispiele für verteilte Ausführung und die Integration von torch.compile.

Der grundlegende Migrationsschritt sieht bewusst vertraut aus:

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

Das ist ein Portabilitätstest, keine Performance-Garantie. Laut PyTorch hängt die maximale TPU-Leistung weiterhin von Modellformen, Sharding und Operationen ab, die zur Hardware passen.

Was vor der Installation bereitstehen muss

TorchTPU ist nur dort sinnvoll einsetzbar, wo eine Google-TPU bereitgestellt wurde. Das Repository empfiehlt, die Notebooks auf einer TPU-VM auszuführen, und nennt eine v6e-Instanz von Google Compute Engine als einfachsten Einstieg. Auf einem Laptop ohne TPU lassen sich die Pakete zwar installieren, eine tatsächliche TPU-Ausführung lässt sich dort aber nicht prüfen.

Die öffentliche README unterstützt Python 3.11+ und führt torch_tpu==0.2.0 als Version auf. Verwende eine isolierte Umgebung, damit der Versuch nicht versehentlich einen CUDA-Build oder andere inkompatible Abhängigkeiten übernimmt.

torch_tpu ohne CUDA installieren

Lege zunächst eine virtuelle Umgebung an und installiere anschließend das CPU-PyTorch-Wheel zusammen mit der festgelegten TorchTPU-Version:

python -m venv .venv
source .venv/bin/activate

pip install torch_tpu==0.2.0 \
  --extra-index-url https://download.pytorch.org/whl/cpu

Wird der zusätzliche Index weggelassen, kann laut README ein CUDA-PyTorch-Wheel installiert werden – genau die falsche Abhängigkeit für eine TPU-VM.

Für die erste Migration empfiehlt sich das veröffentlichte Wheel. Ein Build aus dem Quellcode setzt Bazel oder Bazelisk sowie die im Repository dokumentierten Optionen --config=wheel_common und, falls erforderlich, --config=no_rbe voraus.

Backend mit einem minimalen PyTorch-Programm prüfen

Führe das folgende Beispiel auf der TPU-VM aus:

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

Ein erfolgreiches Ergebnis mit einem tpu-Gerät bestätigt, dass Tensor-Allokation und Matrixmultiplikation das TPU-Backend erreicht haben. Ist das Gerät nicht verfügbar, solltest du zuerst Laufzeitumgebung und Paketinstallation der VM prüfen, bevor du den Modellcode anfasst.

TorchTPU unterstützt zur Fehlersuche die Eager-Ausführung und für kompilierte Ausführung torch.compile. In der PyTorch-Ankündigung wird berichtet, dass der Modus DeferAndFuse in Tests des Projekts mindestens 50% schneller als der strikte Eager-Modus sein kann. Das ist als Aussage des Projekts zu verstehen, nicht als allgemeingültiger Benchmark.

Eine echte PyTorch-Anwendung portieren

Beginne mit der kleinstmöglichen Änderung. Eine TPU-spezifische Konfiguration für verteilte Ausführung kommt erst hinzu, wenn der Pfad auf einem einzelnen Gerät funktioniert.

1. Ersetze die Geräteauswahl durch torch.device("tpu") oder verwende, sofern unterstützt, torch.accelerator.current_accelerator().
2. Verschiebe Modell und Eingabetensoren auf dasselbe Gerät.
3. Führe zunächst einen kurzen Batch aus, bevor du Kompilierung oder verteilte Ausführung aktivierst.
4. Halte Batch-Größen und Tensorformen während der Messung konstant, denn Änderungen an den Formen können eine erneute Kompilierung auslösen.
5. Teste den repräsentativen Modellpfad und das geplante Verteilungsmuster, bevor du eine umfangreichere Migration startest.

Ein verteilter CUDA-Code verwendet häufig NCCL. Das veröffentlichte TorchTPU-Beispiel setzt stattdessen auf das TPU-Process-Group-Backend:

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

Der genaue Startbefehl für die verteilte Ausführung hängt von TPU-Umgebung und Workload ab. Den Device-String zu ändern, konfiguriert weder Sharding noch Prozessplatzierung oder Eingabepipelines.

Der erste Lauf wird voraussichtlich anders ausfallen als spätere Durchläufe. TorchTPU kann laut Dokumentation wiederholte Graphen im Kompilierungs-Cache wiederverwenden. Miss deshalb die Aufwärmphase getrennt vom Durchsatz im stabilen Betrieb.

Wo TorchTPU schon weit ist – und wo noch Nacharbeit nötig bleibt

Google dokumentiert TorchTPU-Integrationen für vLLM, SGLang, TorchTitan und Tools aus dem Hugging-Face-Ökosystem. Die Ankündigung zur Inferenz beschreibt native TPU-Pfade für vLLM und SGLang, darunter die Übertragung des KV-Caches und Funktionen für das Serving. In Googles Cloud-Ankündigung werden außerdem ein MFU-Ergebnis von 52% für Qwen 3 8B mit TorchTitan sowie Angaben zum produktiven Einsatz genannt.

In der Praxis sind die Grenzen weniger glatt als es manche Überschrift vermuten lässt:

Nach dieser Regel solltest du TorchTPU einsetzen

Ein Test mit TorchTPU lohnt sich jetzt, wenn du bereits Zugriff auf Google-TPUs hast, PyTorch-Modelle betreibst und einen gemeinsamen Framework-Pfad für Training oder Inferenz suchst. Der überzeugendste frühe Anwendungsfall ist ein Team, das Tensorformen profiliert und optimiert, statt von einem Einzeiler-Port automatisch eine GPU-ähnliche Effizienz zu erwarten.

Mit der Migration eines kritischen Workloads solltest du warten, wenn du eine ausgereifte Kompatibilitätsmatrix, vorhersehbares Verhalten beim Kaltstart oder eine für deinen Stack wichtige Funktion brauchst, die TorchTPU derzeit noch nicht bereitstellt. Führe die Umstellung schrittweise durch: ein Modell, ein TPU-Slice und ein festes Evaluationsset. Miss Installation, erste Kompilierung, stabilen Durchsatz und Ergebniskorrektheit jeweils getrennt.

TorchTPU-FAQ

Ist TorchTPU dasselbe wie PyTorch/XLA?

Nein. TorchTPU ist ein neueres natives PyTorch-Backend, das auf Eager Execution und torch.compile ausgelegt ist. PyTorch/XLA ist der ältere Integrationspfad für TPUs. Vorhandener XLA-Code erfordert daher möglicherweise eine eigene Migrationsprüfung.

Kann ich TorchTPU auf einem Laptop ohne TPU installieren?

Die Pakete lassen sich möglicherweise installieren. Ohne kompatible Google-TPU-Infrastruktur kannst du die TPU-Ausführung jedoch nicht validieren. Das offizielle Repository geht für den praktischen Einsatz von einer TPU-VM aus.

Warum war mein erster TorchTPU-Schritt so langsam?

Bei der ersten Ausführung kann ein Graph kompiliert und der Cache gefüllt werden. Miss spätere Iterationen separat und halte die Tensorformen konstant, während du erneute Kompilierungen untersuchst.

Garantiert der Wechsel von cuda zu tpu dieselbe Performance?

Nein. Der Austausch des Geräts ist zunächst ein Port auf Korrektheitsebene. Vor einem Performance-Vergleich solltest du Tensorformen, Sharding, Operatorabdeckung und Aufwärmphase messen.

Ist TorchTPU bereit für den Produktiveinsatz?

Google beschreibt produktive Anwendungen für Inferenz und Training. Die Komponenten und Funktionen des öffentlichen Ökosystems befinden sich jedoch auf unterschiedlichen Reifestufen. Teste daher das konkrete Modell, die Serving-Engine und das Verteilungsmuster, bevor du dich festlegst.

Installiere das festgelegte Wheel auf einer TPU-VM, führe die 2×2-Prüfung aus und vergleiche einen aufgewärmten Workload mit deinem aktuellen GPU-Pfad. So lassen sich Code-Portabilität, Compiler-Aufwärmung, Operatorabdeckung und die Wirtschaftlichkeit der Hardware getrennt bewerten.