Google und Meta bringen TorchTitan mit PyTorch-nativen Trainingsrezepten auf TPUs

Ein natives TPU-Backend für PyTorch

Wer PyTorch-Modelle auf Cloud TPUs trainieren wollte, hatte bisher laut dem Google Developers Blog im Wesentlichen zwei Möglichkeiten: die Codebasis in JAX neu schreiben oder mit dem veralteten torch_xla arbeiten. Das habe sich mit dem Start von TorchTPU geändert, so Google. Es stellt Cloud TPUs ein natives Backend bereit, torch.device("tpu"), das in die Kernschnittstelle von PyTorch integriert ist. Damit ein Trainingsskript TPUs nutzt, sind nur noch minimale Änderungen am Code zur Geräteinitialisierung nötig.

Geräteunterstützung allein reiche nicht aus, argumentiert der Beitrag, denn Entwickler bräuchten auch leistungsstarke, produktionsreife Rezepte für verteiltes Training. Diese Rolle übernimmt TorchTitan TPU, ein für Google Cloud TPUs angepasster Fork von Metas TorchTitan. Er ist unter github.com/google-pytorch/torchtitan verfügbar. Google und Meta haben die TPU-nativen Execution Passes in TorchTitan gemeinsam entwickelt, sodass Entwicklungsteams weder eigene TPU-Sharding-Strategien erfinden noch Low-Level-XLA-Kommunikationspasses von Hand schreiben müssen.

Rezepte, entkoppelte Modelle und austauschbare Kernel

TorchTitan bietet Rezepte für das Pre-Training großer Dense-Modelle wie Llama 3 70B/405B und Qwen 3 8B/32B, für Mixture-of-Experts-Modelle wie Qwen 3 235B MoE sowie für Diffusion Transformer wie FLUX.1 12B. Außerdem dient es als Trainer in Post-Training-Pipelines wie VeRL für Group Relative Policy Optimization (GRPO) und für das Alignment per Reinforcement Learning. In diesen Workflows übernimmt es das Parameter-Update und synchronisiert die Gewichte mit Serving-Engines wie vLLM-TorchTPU.

Modelldefinitionen werden in reinem PyTorch für ein einzelnes Gerät geschrieben, ganz ohne CUDA-spezifischen Code. Die Collectives ergänzt eine Framework-Schicht wie FSDP, deklarative Sharding-Pläne wie Tensor Parallelism oder Backend-Plugins wie MoE-Layer. TorchTPU unterstützt die C10D-Collective-APIs und die PyTorch-Abstraktion DeviceMesh.

Der Konfigurationsbaum, aus dem ein Modell aufgebaut wird, lässt sich vor der Konstruktion umschreiben. Dadurch können die TPU-Rezepte TPU-freundliche Bausteine einsetzen, ohne die Modelldefinition oder die gemeinsame Trainer-Schleife anzufassen. Google nennt drei Ersetzungen:

• Splash Attention, ein Pallas-Kernel, der kausale Masken, Sliding-Window-Masken und Masken für gepackte Dokumente verarbeitet, ohne die Attention-Matrix zu materialisieren.
• Ein fusionierter Pallas-Kernel für die finale Projektion und den Cross-Entropy-Loss, der den vollständigen Logits-Tensor nie in den HBM schreibt.
• Gruppierte Pallas-Kernel für die Matmuls der MoE-Experten.

Jeder Kernel ist eine JAX-Funktion, die der Decorator torch.tpu.pallas.jax_op in eine PyTorch-Custom-Op verwandelt, und jede Ersetzung ist ein Konfigurations-Flag. Die Kernel können aus tokamax stammen, der OpenXLA-Kernelbibliothek für TPU und GPU.

SimpleFSDP und Graph Trainer

Der Beitrag erklärt, dass FSDP2 auf das GPU-Ausführungsmodell zugeschnitten ist: Collectives werden aus Modul-Hooks ausgelöst und mithilfe von CUDA-Streams überlappt. Ein Compiler kann das nicht erfassen, weshalb auf TPU eine torch.compile-Region an jedem FSDP-Hook abbricht und sich kein Graph für den gesamten Schritt aufzeichnen lässt. TorchTPU stellt lediglich Stream- und Event-Objekte ohne Funktion bereit, damit FSDP2 überhaupt läuft.

SimpleFSDP wählt stattdessen die Perspektive des Compilers. Das Parameter-Sharding wird zu einer DTensor-Parametrisierung, sodass All-Gather und Reduce-Scatter gewöhnliche Knoten im aufgezeichneten Graphen sind. XLA kann dann All-Gathers vor den Matmuls einplanen, die sie benötigen, und Reduce-Scatters mit der Backward-Berechnung überlappen. Mixed Precision wird über eine Policy auf demselben Mechanismus gesteuert, und hybrides Sharding für Multislice nutzt ein zweidimensionales Mesh. Bei FLUX brachte das TorchTitan laut Google auf 99 % der Leistung von JAX MaxDiffusion.

Graph Trainer geht weiter und erfasst den vollständigen Trainingsschritt (Forward, Loss und Backward) mit dem make_fx-Tracer von PyTorch als einen einzigen FX-Graphen, sodass der Schritt als ein einziges XLA-Programm läuft. Laut dem Beitrag ist PyTorch-Training auf TPU damit konzeptionell dasselbe wie JAX. Außerdem ermöglicht es Scan über Layer und einen Graph Splitter, der Aktivierungen zwischen separat kompilierten Forward- und Backward-Programmen in den Host-Speicher auslagert. Laut Google erreicht dieser Pfad bei Qwen3 235B mehr als 29 % MFU. Mit einem funktionalen AdamW-Update, das in den Schritt eingebunden ist, kommt Qwen3-8B auf 51 bis 52 % MFU.

Ein Playbook aus Konfigurationsänderungen

Google hat zudem ein Playbook zur Optimierung von TPU-Modellen veröffentlicht, das im TorchTitan-TPU-Optimierungsleitfaden beschrieben ist und dessen Hebel fast alle bereits in die Rezepte eingebunden sind. Laut dem Beitrag ist es bei den meisten Modellen "eine Abfolge von Konfigurationsänderungen an einem bestehenden Rezept statt neuer Code", um auf TPU eine gute MFU zu erreichen. Die Reihenfolge ist wichtig, da jeder Schritt verändert, was der nächste misst.

Der erste Schritt ist die Wahl einer Baseline: Eager, Compile oder Graph Trainer. Der Eager-Modus führt PyTorch-Modelle auf TPU praktisch ohne Änderungen aus, was sich für Inbetriebnahme und Debugging eines Modells eignet, aber keine Leistungskonfiguration ist, weil jede Op ein eigenes Programm darstellt. Der Deferred-and-Fused-Eager-Modus von TorchTPU holt einen Teil davon zurück, indem XLA über Op-Grenzen hinweg fusionieren kann. Für PyTorch-Teams lautet die praktische Botschaft, dass TPU-Training nicht mehr mit einem JAX-Rewrite beginnen muss. Der vorliegende Text bricht vor den übrigen Schritten des Playbooks ab.