Google gör TorchTPU till öppen källkod och ger inbyggt PyTorch-stöd på sina TPU:er

Google meddelade den 20 oktober 2026 på PyTorch Conference 2026 att TorchTPU, företagets inbyggda PyTorch-stapel för Tensor Processing Units (TPU:er), nu finns som öppen källkod på GitHub, meddelade företaget på Google Cloud-bloggen. Programvaran har utvecklats tillsammans med PyTorch-teamet på Meta och är utformad för att köra befintliga PyTorch-arbetsbelastningar på TPU:er med minimala kodändringar.

Med den här lanseringen blir TPU:er ett inbyggt alternativ för PyTorch, vid sidan av GPU:er och CPU:er, med stöd för träning, inferens och driftsättning genom verktyg som vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face och LiteRT, uppger Google. Utvecklare kan felsöka modeller i eager-läge och sedan uppnå prestanda på produktionsnivå med torch.compile för att optimera TPU-utnyttjandet och minska minnesåtgången.

Jana van Greunen, chef för PyTorch Engineering på Meta, sade i tillkännagivandet att TorchTPU gör det möjligt för Meta att flexibelt samutveckla hårdvara med företagets centrala team för maskininlärningssystem, samtidigt som stödet för stora språkmodeller (LLM), annonseringssystem och rekommendationssystem säkerställs. ”Det gör det möjligt för oss att driftsätta våra PyTorch-arbetsbelastningar på Googles TPU:er, ge direkt återkoppling till utvecklingsteamen och bygga den flexibilitet i ekosystemet som alla behöver”, sade hon.

Integrationer för Hugging och modellservering

Google samarbetade med Hugging Face för att öppna viktmodeller på Hugging Face Hub ska köras direkt på TorchTPU. I Transformers v5 driftsätter man en modell på TPU:er genom att ange device_map="tpu", och Diffusers-biblioteket körs direkt på TPU:er med stöd för eager-läge, torch.compile och tensorparallellism över flera acceleratorer. Google uppger att företaget också byggt särskilda CI/CD-pipelines för kontinuerlig nattlig testning av Hugging Face Transformers-sviten på TPU-infrastruktur. Jeff Boudier, produktchef på Hugging Face, sade att samarbetet omsätter plattformens åtagande att erbjuda valfrihet när det gäller hårdvara i praktiken.

Den inbyggda PyTorch-backenden för vLLM på TorchTPU har utvecklats i samarbete med Inferact för modellservering. Google uppger att den stöder kontinuerlig batchning, kvantisering och parallellism över flera chip på TPU-hårdvara med anpassade Pallas-kärnor, samtidigt som de befintliga worker-modellerna, schemaläggarna och API-servrarna bevaras så att befintliga inferenssystem inte behöver ändras. Simon Mo, vd och medgrundare av Inferact, sade att integrationen ger TPU:er förstklassigt stöd för vLLM.

En inbyggd PyTorch-backend för SGLang, utvecklad tillsammans med RadixArk, är avsedd för strukturerade arbetsbelastningar och agentbaserade arbetsbelastningar. Enligt tillkännagivandet använder SGLang RadixAttention för att återanvända KV-cachen mellan förfrågningar, så att upprepade promptprefix inte behöver beräknas på nytt. TorchTPU-versionen använder dessutom en backend för överföring av separerade arbetsbelastningar, som också fungerar som I/O-backend för HiCache. RadixArks teknikchef och medgrundare Banghua Zhu sade att TorchTPU har påskyndat företagets arbete med att införa stöd för SGLang på olika chip.

Google lanserar också llm-d 1.0 tillsammans med Red Hat, IBM och en gemenskap med över 500 bidragsgivare och 25 organisationer, som ett Sandbox-projekt inom Cloud Native Computing Foundation. Tillkännagivandet beskriver llm-d som en länk mellan PyTorch-ekosystemet för modellservering och produktionsdriftsättning baserad på Kubernetes: routern använder prefixcachemedveten routning för att skicka varje förfrågan till den replik som redan har dess kontext, och separerar den beräkningsintensiva förifyllningen från den minnesintensiva avkodningen. Version 1.0 lägger till flödeskontroll för flera klienter, hög tillgänglighet för routern, stabila v1-API:er och sessionsmedveten routning för agenter. På Google Kubernetes Engine kan en gemensam llm-d-slutpunkt köras på antingen GPU:er eller TPU:er.

Prof. Sashikumaar Ganesan, grundare av Zenteiq, sade i tillkännagivandet att GKE Inference Gateway, som drivs av llm-d, gav över 4× högre sammanlagd inferensgenomströmning för företagets syntetiska datagenerering. Den dagliga kapaciteten för genererade token ökade från 1,47 miljarder till 5,97 miljarder token per 1 000 TPU v6e-acceleratorer.

Träning och orkestrering

TorchTitan, PyTorchs inbyggda referensimplementation för distribuerad träning, innehåller nu en integration med TorchTPU som enligt Google håller produktionsnivå och bygger på torch.compile, TorchTitans SimpleFSDP, specialiserade Pallas-kärnor och optimeringar på grafnivå. TorchTitan tillhandahåller TorchTPU-recept för förträning av täta modeller som Llama 3 70B och 405B, Qwen 3 8B och 32B samt MuseGlimmer 30B, liksom blandexpertsmodellen Qwen 3 235B och diffusionstransformern FLUX.1 12B. Google uppger att integrationen har validerats med upp till 8 000 Ironwood-TPU:er (TPU v7x) och att träning med TorchTPU via TorchTitan nådde upp till 52 % utnyttjande av modellens FLOPS för Qwen 3 8B, med nära linjär skalning över upp till 8 192 Ironwood-TPU:er.

Google bidrog också med en TPU-backend och ett referensrecept från början till slut till VeRL, ett ramverk för förstärkningsinlärning. GRPO-receptet, som bygger på en arkitektur med separerade tränar- och samplarenheter och har demonstrerats på dataset för matematiskt resonemang, är avsett att låta utvecklare köra VeRL-arbetsbelastningar på TPU:er utan att behöva ändra sin befintliga hårdvaruoberoende modell- och algoritmkod.

Ray ger inbyggt TorchTPU-stöd för att samordna PyTorch-databehandling, träning, förstärkningsinlärning och modellservering på CPU:er, GPU:er och TPU:er i ett och samma arbetsflöde. Edward Oakes, chefsarkitekt på Anyscale, sade att företagen arbetade för att göra TPU:er till acceleratorer av första klass i Ray och tillade: ”TorchTPU integrerar Google TPU:er direkt i den branschstandardiserade programvarustacken.”

Tidigare arkitektur och tillgänglighet

TorchTPU presenterades offentligt den 7 april 2026 i ett blogginlägg på Google Developers Blog som beskrev en ”Eager First”-design byggd på PyTorchs PrivateUse1-gränssnitt, med lägena Debug, Strict och Fused för eager-körning, som alla använder en gemensam kompileringscache. Google uppgav då att Fused Eager ger en prestandaökning på 50 % till 100+ % jämfört med Strict Eager, och att torch.compile i TorchTPU fångar FX-grafer med Torch Dynamo och sänker dem via XLA och StableHLO, med anpassade kärnor skrivna i Pallas och JAX. I blogginlägget stod det att programvarustacken har inbyggt stöd för Distributed Data Parallel, FSDPv2 och DTensor och är utformad för att hantera avvikande exekvering mellan ranker. Där angavs också lanseringen av ett offentligt GitHub-arkiv som en av TorchTPU:s färdplanerade punkter för 2026.

Arkivet för torch_tpu är nu offentligt under GitHub-organisationen google-pytorch och licensierat under Apache 2.0. I README-filen listas anpassade ATen-kärnor, en kompileringscache för XLA-beräkningar, stöd för distribuerad träning med data- och tensorparallellism, integrering med torch.compile, verktyg för felsökning och prestandamätning samt exempel för Llama, Qwen, ResNet, DLRM, minGPT och UNet. Projektet stöder Python 3.11 och senare, och paketet torch-tpu 0.2.0 publicerades på PyPI den 19 oktober 2026, med Linux-hjul för x86-64 och CPython 3.11 till och med 3.14. En användarhandbok finns på torchtpu.readthedocs.io, och i README-filen står det att projektets interaktiva notebookar måste köras från en TPU VM, där en v6e GCE-instans är det enklaste alternativet.

I samma tillkännagivande uppgav Google att LiteRT, tidigare TensorFlow Lite, ansluter sig till PyTorch-ekosystemet: biblioteket LiteRT-Torch konverterar en PyTorch-modell med ett enda litert-torch.convert() anrop, och konverterade modeller körs på Android, iOS, macOS, Windows, Linux, webben och IoT-enheter. En uppsättning TPU-agenter, avsedd att automatisera introduktion, diagnostik och finjustering av inferensprestanda för PyTorch-modeller, finns i en tidig förhandsversion på GitHub.

Tillkännagivandet hänvisar utvecklare till Google Clouds TPU-dokumentation för att komma igång med att köra PyTorch på TPU:er.