Google hat als Open Source veröffentlicht TorchTPU, einen nativen PyTorch-Backend für Cloud TPUs. Er wurde gemeinsam mit dem PyTorch-Team bei Meta entwickelt und auf der PyTorch Conference 2026 am 20. Oktober 2026 angekündigt. Vorhandener PyTorch-Code richtet sich über ein Standard-Target an eine TPU torch.device("tpu"). TPUs stehen nun neben GPUs und CPUs als natives PyTorch-Device. Das Google-Team gibt an, dass TorchTPU PyTorch/XLA ersetzt und empfiehlt es für alle neuen Workloads.
TL;DR
- Größe: Ein Software-Stack, kein Modell. 4 offene Repos: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan. TorchTPU ist v0.2.0 auf PyPI, Python 3.11+, Apache 2.0.
- Läuft auf: Google Cloud TPUs. Getestet auf Generationen ab v5e, skaliert auf 8,192 Ironwood (TPU7x)-Chips. Serving unterstützt FP8-, MXFP4- und NVFP4-Checkpoints.
- Leistung: Eager-Modus zum Debuggen,
torch.compilefür Throughput auf Produktionsniveau. - Bestwert: 52% MFU beim Training von Qwen 3 8B mit TorchTitan auf Ironwood (selbst gemeldet).
- Schlechtester Wert: Out-of-the-box-Eager benötigte in einem Fine-Tuning-Fall 11.49s pro Schritt, gegenüber 2.78s voll optimiert.
Was ist TorchTPU?
TorchTPU ist ein Open-Source-PyTorch-Backend von Google und Meta, das PyTorch-Modelle nativ auf TPUs ausführt. Es bindet sich über PyTorchs PrivateUse1-Interfaceein. Tensoren sind gewöhnliche PyTorch-Tensoren, ohne Subclasses oder Wrapper. Es läuft mit einem Prozess pro Device und erhält Tensor-Views sowie In-Place-Mutationen.
Der Vorgänger, PyTorch/XLA, zeichnete Lazy-Tensor-Graphen auf. Entwickler fügten häufig Code hinzu, um kompilierbare Regionen zu markieren. PyTorch/XLA auch unterstützte nur reinen SPMD-Code. TorchTPU beherrscht außerdem divergente MPMD-Ausführung.
Wie führt TorchTPU PyTorch-Code aus?
TorchTPU bietet 3 Eager-Modi:
- Strict (Standard): verteilt jeden Operator asynchron mit Standard-PyTorch-Semantik.
- Debug: läuft synchron mit Tensor-Grenzwertprüfungen.
- DeferAndFuse: verschmilzt Folgen von Eager-Operatoren, um HBM-Traffic zu reduzieren. Das TorchTPU-Papier besagt, dass dies 50% oder mehr schneller als Strict sein kann.
Für Höchstgeschwindigkeit, torch.compile(backend="tpu") erfasst den Graphen mit TorchDynamo. Es wird über StableHLO zu OpenXLA abgesenkt, das TPU-Maschinencode erzeugt. Kompilierte ausführbare Dateien werden prozess-, host- und neustartübergreifend zwischengespeichert.
Benutzerdefinierte Pallas-Kernel sind als normale PyTorch-Funktionen aufrufbar. Vorgefertigte Tokamax -Kernel und Helion -Tile-Kernel funktionieren ebenfalls.
Die Verteilung nutzt die tpu_dist Prozessgruppe. Sie unterstützt DDP, FSDP2 und DTensor sowie ein topologiebewusstes Mesh für TPU-Interconnects.
Wie schnell ist TorchTPU?
Alle folgenden Zahlen wurden von Google und Partnern selbst berichtet.
- Trainingsskalierung: TorchTitan erreichte bis zu 52% MFU bei Qwen 3 8B. Google berichtet von nahezu linearer Skalierung bis zu 8,192 Ironwood TPUs.
- Tuning-Leiter: In einem TorchTitan-Fine-Tuning-Fall sank die Schrittzeit von 11.49s auf 2.78s. DeferAndFuse allein lieferte 1.15x. Vollständig
torch.compilelieferte 3.80x. - Modellabdeckung: Google testet kontinuierlich etwa 5,000 Hugging Face Modelle im eager- und kompilierten Modus.
- Fleet Serving: Zenteiq steigerte die tägliche Ausgabe von 1.47B auf 5.97B Token pro 1,000 TPU v6e Chips. Dieser Gewinn kam vom llm-d Routing auf GKE, nicht von TorchTPU allein.
Welche Integrationen werden mit TorchTPU ausgeliefert?
- Hugging Face: Transformers v5 wird mit
device_map="tpu"auf TPU bereitgestellt. Diffusers unterstützt eager,torch.compileund Tensorparallelität. - vLLM-TorchTPU: gebaut mit Inferact. Behält die Upstream-Scheduler und API-Server bei, mit benutzerdefinierten Pallas-Kernels.
- SGLang-TorchTPU: mit RadixArk gebaut. Behält RadixAttention-Prefix-Reuse bei.
- Serving-Funktionen: FP8-KV-Caches, FP8-, MXFP4- und NVFP4-Checkpoints. Speculative Decoding umfasst MTP und EAGLE-3.
- TorchTitan: Rezepte für Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE und FLUX.1 12B.
- Außerdem: llm-d 1.0, Ray, TorchRec, VeRL (Vorschau) und LiteRT-Torch für On-Device-Export.
Wie schneidet TorchTPU im Vergleich zu PyTorch/XLA, TorchNeuron und JAX ab?
| Merkmal | TorchTPU (Repo) | PyTorch/XLA (Repo) | TorchNeuron (AWS-Dokumentation) | JAX on TPU (Google-Dokumentation) |
|---|---|---|---|---|
| Anbieter | Google + Meta | Google / PyTorch | AWS | |
| Hardware | Cloud TPU (v5e und später getestet) | Cloud TPU | Trainium 2, Trainium 3 | Cloud TPU |
| PyTorch-Gerät | tpu | xla | neuron | Kein PyTorch |
| Backend-Mechanismus | PrivateUse1 | Lazy-Tensor-Graph-Aufzeichnung | PrivateUse1 | Nicht anwendbar |
| Standard-Ausführung | Eager (Strict, Debug, DeferAndFuse) | Lazy Tensors | Eager plus torch.compile | Nicht offengelegt |
| Compile-Pfad | TorchDynamo zu StableHLO zu XLA | XLA | Neuron Compiler | XLA |
| Distributed APIs | DDP, FSDP2, DTensor, SPMD und MPMD | Nur reines SPMD | FSDP, DTensor, DDP, TP | Nicht offengelegt |
| Eigene Kernels | Pallas, Tokamax, Helion | Nicht offengelegt | NKI | Pallas |
| vLLM / SGLang Serving | vLLM-TorchTPU, SGLang-TorchTPU | Nicht offengelegt | vLLM Neuron | vLLM-JAX, SGLang-JAX |
| Status | Open Source, v0.2.0 | Wird durch TorchTPU ersetzt | Öffentliche Beta (Neuron 2.33) | Gepflegt |
| Lizenz | Apache 2.0 | Nicht offengelegt | Nicht offengelegt | Nicht offengelegt |
Wie installiert man TorchTPU?
Installieren Sie das Paket auf einer TPU-VM und wählen Sie dann das tpu Gerät:
# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)Die README des Repos pinnt torch_tpu==0.2.0 mit dem CPU-PyTorch-Wheel-Index. Ohne diese lädt pip die CUDA-Version von torchherunter. Vollständige Anleitungen finden sich unter torchtpu.readthedocs.io.
Wichtigste Erkenntnisse
- TorchTPU ist Open Source unter Apache 2.0 und wurde gemeinsam mit Meta entwickelt.
- Es ersetzt PyTorch/XLA durch eine Eager-First-Ausführung auf Basis von PrivateUse1.
torch.compile(backend="tpu")senkt die Kosten durch StableHLO und XLA.- TorchTitan erreichte 52% MFU auf Qwen 3 8B und skalierte auf 8,192 Ironwood-Chips.
- vLLM, SGLang und Hugging Face laufen nun über Standard-PyTorch auf TPU.
Sehen Sie sich das GitHub-Repo, Paper, Dokumentation und die Google-Cloud-Ankündigungan. Alle Credits gehen an die Forscher dieses Projekts. Folgen Sie uns gerne auch auf Twitter und vergessen Sie nicht, unserem 150k+ML SubReddit beizutreten und unseren Newsletterzu abonnieren. Moment! Sind Sie auf Telegram? jetzt können Sie uns auch auf Telegram beitreten.
Der Beitrag Google macht TorchTPU als Open Source verfügbar, ein natives PyTorch-Backend für Cloud-TPUs erschien zuerst auf MarkTechPost.