Google maakt TorchTPU opensource en brengt PyTorch rechtstreeks naar zijn TPU’s

Google maakte op 20 oktober 2026 tijdens PyTorch Conference 2026 bekend dat TorchTPU, Googles native PyTorch-stack voor Tensor Processing Units (TPU’s), nu als opensourcesoftware beschikbaar is op GitHub, maakte het bedrijf bekend op de Google Cloud-blog. De software is samen met het PyTorch-team van Meta ontwikkeld en is bedoeld om bestaande PyTorch-workloads op TPU’s uit te voeren met minimale wijzigingen in de code.

Met deze release zijn TPU’s naast GPU’s en CPU’s een native optie voor PyTorch, aldus Google. De stack bestrijkt training, serving en implementatie en omvat tools als vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face en LiteRT. Ontwikkelaars kunnen modellen debuggen in eager mode en vervolgens met torch.compile overstappen op prestaties van productieniveau, om het TPU-gebruik te optimaliseren en het geheugengebruik te verminderen.

Jana van Greunen, directeur PyTorch Engineering bij Meta, zei in de aankondiging dat TorchTPU Meta in staat stelt hardware flexibel samen met het centrale machinelearningteam te ontwikkelen, terwijl het bedrijf tegelijkertijd workloads voor grote LLM’s, advertentie- en aanbevelingssystemen ondersteunt. “Zo kunnen we onze PyTorch-workloads op Google TPU’s uitvoeren, rechtstreeks feedback geven aan de engineeringteams en de flexibiliteit binnen het ecosysteem ontwikkelen die iedereen nodig heeft”, zei ze.

Integraties voor Hugging en serving

Google ging een samenwerking aan met Hugging Face, zodat open-gewichtsmodellen op de Hugging Face Hub native op TorchTPU draaien. In Transformers v5 implementeer je een model op TPU’s door device_map="tpu", en de Diffusers-bibliotheek draait native op TPU’s en ondersteunt eager mode, torch.compile en tensorparallelisme over meerdere accelerators. Google zei ook speciale CI/CD-pijplijnen te hebben gebouwd om de Transformers-suite van Hugging Face elke nacht continu op TPU-infrastructuur te testen. Jeff Boudier, VP Product bij Hugging Face, zei dat de samenwerking de toezegging van het platform om hardwarekeuze te bieden in de praktijk brengt.

De native PyTorch-vLLM-backend op TorchTPU is voor serving ontwikkeld in samenwerking met Inferact. Volgens Google ondersteunt deze backend continuous batching, kwantisering en parallelle verwerking over meerdere chips op TPU-hardware met behulp van aangepaste Pallas-kernels. Daarbij blijven de upstream-workermodellen, schedulers en API-servers behouden, zodat bestaande inferentiestacks niet hoeven te worden aangepast. Simon Mo, CEO en medeoprichter van Inferact, zei dat de integratie TPU’s volwaardige vLLM-ondersteuning biedt.

Een native PyTorch-backend voor SGLang, ontwikkeld in samenwerking met RadixArk, is gericht op gestructureerde workloads en workloads met AI-agenten. Volgens de aankondiging gebruikt SGLang RadixAttention om de KV-cache tussen verzoeken te hergebruiken, zodat terugkerende promptprefixen niet opnieuw hoeven te worden berekend. De TorchTPU-versie gebruikt daarnaast een backend voor gedisaggregeerde overdracht, die ook dienstdoet als I/O-backend voor HiCache. Banghua Zhu, CTO en medeoprichter van RadixArk, zei dat TorchTPU het werk van het bedrijf heeft versneld terwijl het SGLang-ondersteuning uitrolt voor verschillende soorten silicium.

Google brengt ook llm-d 1.0 uit met Red Hat, IBM en een gemeenschap van meer dan 500 bijdragers en 25 organisaties, als Sandbox-project van de Cloud Native Computing Foundation. De aankondiging beschrijft llm-d als een verbinding tussen het PyTorch-servingecosysteem en productie-implementaties op basis van Kubernetes: de router gebruikt routering die rekening houdt met prefixcaches, zodat elk verzoek naar de replica gaat die de bijbehorende context al bevat. Ook worden de rekenintensieve prefill-fase en de geheugenintensieve decode-fase van elkaar gescheiden. Versie 1.0 voegt flowcontrol voor meerdere tenants, hoge beschikbaarheid voor de router, stabiele v1-API’s en sessiebewuste routering voor agents toe. Op Google Kubernetes Engine kan één gezamenlijk llm-d-eindpunt draaien op GPU’s of TPU’s.

Prof. Sashikumaar Ganesan, oprichter van Zenteiq, zei in de aankondiging dat GKE Inference Gateway, aangedreven door llm-d, voor het genereren van synthetische data bij het bedrijf een meer dan 4× hogere totale inferentiedoorvoer opleverde. Daardoor steeg de dagelijkse capaciteit voor het genereren van tokens van 1,47 miljard naar 5,97 miljard tokens per 1.000 TPU v6e-accelerators.

Training en orkestratie

TorchTitan, de native PyTorch-referentie-implementatie voor gedistribueerde training, bevat nu een integratie met TorchTPU van productiekwaliteit, die volgens Google is gebouwd op torch.compile, TorchTitans SimpleFSDP, gespecialiseerde Pallas-kernels en optimalisaties op graafniveau. TorchTitan biedt TorchTPU-recepten voor het voortrainen van dichte modellen, zoals Llama 3 70B en 405B, Qwen 3 8B en 32B en MuseGlimmer 30B, en van het mixture-of-experts-model Qwen 3 235B en de diffusietransformer FLUX.1 12B. Google zei dat de integratie is gevalideerd op maximaal 8.000 Ironwood-TPU’s (TPU v7x) en dat training met TorchTPU via TorchTitan op Qwen 3 8B een Model Flops Utilization van maximaal 52% heeft bereikt, met een vrijwel lineaire schaalbaarheid op maximaal 8.192 Ironwood-TPU’s.

Google heeft ook een TPU-backend en een end-to-end-referentierecept bijgedragen aan VeRL, een framework voor reinforcement learning. Het GRPO-recept, gebouwd op een architectuur met een gescheiden trainer en sampler en gedemonstreerd met datasets voor wiskundig redeneren, moet ontwikkelaars in staat stellen VeRL-workloads op TPU’s uit te voeren met behoud van hun bestaande hardwareonafhankelijke code voor modellen en algoritmen.

Ray voegt native ondersteuning voor TorchTPU toe, zodat PyTorch-workflows voor gegevensverwerking, training, versterkend leren en serving kunnen worden georkestreerd op CPU’s, GPU’s en TPU’s. Edward Oakes, hoofdarchitect bij Anyscale, zei dat de bedrijven eraan hebben gewerkt om TPU’s binnen Ray als volwaardige accelerators te ondersteunen. Hij voegde eraan toe: „TorchTPU integreert Google TPU’s rechtstreeks in die industriestandaardstack.”

Eerdere architectuur en beschikbaarheid

TorchTPU werd op 7 april 2026 publiekelijk geïntroduceerd in een blogbericht van Google Developers. Daarin werd een ontwerp met de naam „Eager First” beschreven, gebaseerd op PyTorchs PrivateUse1-interface, met de eager-modi Debug, Strict en Fused, die gebruikmaken van een gedeelde compilatiecache. Google zei destijds dat Fused Eager een prestatieverbetering van 50% tot 100% of meer oplevert ten opzichte van Strict Eager, en dat torch.compile op TorchTPU FX-grafieken vastlegt met Torch Dynamo en deze via XLA en StableHLO omzet, met aangepaste kernels die zijn geschreven in Pallas en JAX. Volgens het bericht ondersteunt de stack standaard Distributed Data Parallel, FSDPv2 en DTensor, en is deze ontworpen om uiteenlopende uitvoering over verschillende ranks aan te kunnen. Ook noemde het bericht de lancering van een openbare GitHub-opslagplaats als een van de roadmapdoelen van TorchTPU voor 2026.

De torch_tpu-opslagplaats is nu openbaar binnen de GitHub-organisatie google-pytorch en valt onder de Apache 2.0-licentie. In de README staan aangepaste ATen-kernels, een compilatiecache voor XLA-berekeningen, ondersteuning voor gedistribueerde training met data- en tensorparallellisme, integratie met torch.compile, hulpprogramma’s voor foutopsporing en benchmarking, en voorbeelden voor Llama, Qwen, ResNet, DLRM, minGPT en UNet. Het project ondersteunt Python 3.11 en hoger. Het pakket torch-tpu 0.2.0 werd op 19 oktober 2026 uitgebracht op PyPI, met Linux x86-64-wheels voor CPython 3.11 tot en met 3.14. Er is een gebruikershandleiding gepubliceerd op torchtpu.readthedocs.io. De README vermeldt dat de interactieve notebooks van het project vanuit een TPU-VM moeten worden uitgevoerd; een v6e GCE-instantie is daarvoor de eenvoudigste optie.

In dezelfde aankondiging zei Google dat LiteRT, voorheen TensorFlow Lite, onderdeel wordt van het PyTorch-ecosysteem: de LiteRT-Torch-bibliotheek converteert een PyTorch-model met één enkele litert-torch.convert() aanroep, en de geconverteerde modellen draaien op Android, iOS, macOS, Windows, Linux, het web en IoT-apparaten. Een suite met TPU-agents die bedoeld is om de ingebruikname, diagnostiek en het afstemmen van inferentieprestaties voor PyTorch-modellen te automatiseren, bevindt zich in een vroege preview op GitHub.

De aankondiging verwijst ontwikkelaars naar de TPU-documentatie van Google Cloud om PyTorch op TPU’s te gaan gebruiken.