Google frigiver TorchTPU som open source og gør PyTorch til en indbygget mulighed på sine TPU’er

Google annoncerede den 20. oktober 2026 på PyTorch Conference 2026, at TorchTPU, virksomhedens indbyggede PyTorch-stak til Tensor Processing Units (TPU’er), nu er tilgængelig som open source på GitHub, oplyste virksomheden på Google Cloud-bloggen. Softwaren er udviklet i fællesskab med PyTorch-teamet hos Meta og er designet til at køre eksisterende PyTorch-arbejdsbelastninger på TPU’er med minimale kodeændringer.

Med lanceringen bliver TPU’er, ifølge Google, en indbygget mulighed for PyTorch på linje med GPU’er og CPU’er. Det gælder træning, afvikling og implementering via værktøjer som vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face og LiteRT. Udviklere kan fejlfinde modeller i eager mode og derefter opnå ydeevne på produktionsniveau med torch.compile, der optimerer udnyttelsen af TPU’er og reducerer hukommelsesforbruget.

Jana van Greunen, direktør for PyTorch Engineering hos Meta, sagde i annonceringen, at TorchTPU gør det muligt for Meta fleksibelt at udvikle hardware i fællesskab med virksomhedens centrale team for maskinlæringssystemer og samtidig understøtte store arbejdsbelastninger fra LLM-, annonce- og anbefalingssystemer. »Det gør det muligt for os at implementere vores PyTorch-arbejdsbelastninger på Google TPU’er, give direkte feedback til udviklingsteamene og opbygge den fleksibilitet i økosystemet, som alle har brug for,« sagde hun.

Integrationer til Hugging og servering

Google indgik et partnerskab med Hugging Face, så modeller med åbne vægte på Hugging Face Hub kan køre direkte på TorchTPU. I Transformers v5 foregår implementering af en model på TPU’er ved at angive device_map="tpu", og Diffusers-biblioteket kører direkte på TPU’er med understøttelse af eager mode, torch.compile og tensorparallelisering på tværs af flere acceleratorer. Google oplyste, at virksomheden også har opbygget særlige CI/CD-pipelines til løbende natlige tests af Hugging Face Transformers-pakken på TPU-infrastruktur. Jeff Boudier, produktchef hos Hugging Face, sagde, at samarbejdet omsætter platformens løfte om valgfrihed med hensyn til hardware til praksis.

Til servering blev den indbyggede PyTorch vLLM-backend på TorchTPU udviklet i samarbejde med Inferact. Google oplyste, at den understøtter kontinuerlig batching, kvantisering og parallelisering på tværs af flere chips på TPU-hardware ved hjælp af tilpassede Pallas-kerner, samtidig med at den bevarer de oprindelige worker-modeller, schedulere og API-servere, så eksisterende inferensstakke ikke kræver ændringer. Simon Mo, administrerende direktør og medstifter af Inferact, sagde, at integrationen giver TPU’er førsteklasses understøttelse af vLLM.

En indbygget PyTorch-backend til SGLang, udviklet sammen med RadixArk, er rettet mod strukturerede arbejdsbelastninger og agentbaserede arbejdsbelastninger. Ifølge annonceringen bruger SGLang RadixAttention til at genbruge KV-cachen på tværs af anmodninger, så gentagne præfikser i prompts ikke skal beregnes igen, og TorchTPU-versionen anvender en backend til overførsel mellem adskilte komponenter, som også fungerer som I/O-backend for HiCache. RadixArks teknologidirektør og medstifter Banghua Zhu sagde, at TorchTPU har fremskyndet virksomhedens arbejde, mens den udruller understøttelse af SGLang på forskellige typer chips.

Google udgiver også llm-d 1.0 sammen med Red Hat, IBM og et fællesskab på mere end 500 bidragydere og 25 organisationer som et Sandbox-projekt under Cloud Native Computing Foundation. Annonceringen beskriver llm-d som et bindeled mellem PyTorch-økosystemet for servering og produktionsudrulning baseret på Kubernetes: Routeren bruger routing med kendskab til præfikscache til at sende hver anmodning til den replika, der allerede indeholder dens kontekst, og adskiller den beregningstunge præfyldning fra den hukommelsestunge dekodning. Version 1.0 tilføjer flowstyring for flere lejere, høj tilgængelighed for routeren, stabile v1-API’er og sessionsbevidst routing for agenter. På Google Kubernetes Engine kan et fælles llm-d-slutpunkt køre på enten GPU’er eller TPU’er.

Prof. Sashikumaar Ganesan, grundlægger af Zenteiq, sagde i annonceringen, at GKE Inference Gateway drevet af llm-d leverede over 4× højere samlet inferensgennemløb til virksomhedens generering af syntetiske data. Det øgede den daglige kapacitet for genererede tokens fra 1,47 milliarder til 5,97 milliarder tokens pr. 1.000 TPU v6e-acceleratorer.

Træning og orkestrering

TorchTitan, PyTorchs indbyggede referenceimplementering til distribueret træning, indeholder nu det, Google beskriver som en TorchTPU-integration i produktionskvalitet, bygget på torch.compile, TorchTitans SimpleFSDP, specialiserede Pallas-kerner og optimeringer på grafniveau. TorchTitan indeholder TorchTPU-opskrifter til fortræning af tætte modeller som Llama 3 70B og 405B, Qwen 3 8B og 32B samt MuseGlimmer 30B, foruden mixture-of-experts-modellen Qwen 3 235B og diffusionstransformeren FLUX.1 12B. Google oplyste, at integrationen er valideret på op til 8.000 Ironwood TPU’er (TPU v7x), og at TorchTPU-træning via TorchTitan nåede op på 52 % Model Flops Utilization for Qwen 3 8B med næsten lineær skalering på tværs af op til 8.192 Ironwood TPU’er.

Google bidrog også med en TPU-backend og en end-to-end-referenceopskrift til VeRL, et framework til forstærkningslæring. GRPO-opskriften, der bygger på en arkitektur med adskilt træner og sampler og er demonstreret på datasæt til matematisk ræsonnement, skal gøre det muligt for udviklere at køre VeRL-arbejdsbelastninger på TPU’er og samtidig bevare deres eksisterende hardwareuafhængige kode til modeller og algoritmer.

Ray tilføjer indbygget TorchTPU-understøttelse, så PyTorch kan håndtere databehandling, træning, forstærkningslæring og serving på tværs af CPU’er, GPU’er og TPU’er i ét samlet workflow. Edward Oakes, chefarkitekt hos Anyscale, sagde, at virksomhederne havde arbejdet på at gøre TPU’er til acceleratorer på lige fod med andre i Ray, og tilføjede: “TorchTPU integrerer Googles TPU’er direkte i denne branchestandardiserede stak.”

Tidligere arkitektur og tilgængelighed

TorchTPU blev offentligt præsenteret den 7. april 2026 i et blogindlæg på Google Developers Blog, der beskrev et “Eager First”-design baseret på PyTorches PrivateUse1-grænseflade med tilstandene Debug, Strict og Fused eager, som benytter en fælles kompileringscache. Google sagde dengang, at Fused Eager giver en ydelsesforbedring på 50 % til over 100 % sammenlignet med Strict Eager, og at torch.compile på TorchTPU indfanger FX-grafer med Torch Dynamo og oversætter dem via XLA og StableHLO ved hjælp af brugerdefinerede kerner skrevet i Pallas og JAX. Ifølge indlægget understøtter stakken Distributed Data Parallel, FSDPv2 og DTensor direkte fra start, og arkitekturen er udformet til at håndtere forskelligartet eksekvering på tværs af rangnumre. Indlægget nævnte også lanceringen af et offentligt GitHub-lager blandt TorchTPU’s planlagte milepæle for 2026.

torch_tpu-lageret er nu offentligt tilgængeligt under GitHub-organisationen google-pytorch og er licenseret under Apache 2.0. README-filen oplister brugerdefinerede ATen-kerner, en kompileringscache til XLA-beregninger, understøttelse af distribueret træning med dataparallelisme og tensorparallelisme, integration med torch.compile, værktøjer til fejlfinding og benchmarking samt eksempler med Llama, Qwen, ResNet, DLRM, minGPT og UNet. Projektet understøtter Python 3.11 og nyere, og pakken torch-tpu 0.2.0 blev udgivet på PyPI den 19. oktober 2026 med Linux x86-64-hjul til CPython 3.11 til og med 3.14. En brugervejledning findes på torchtpu.readthedocs.io, og README-filen bemærker, at projektets interaktive notesbøger skal køres fra en TPU-VM, hvor en v6e GCE-instans er den nemmeste mulighed.

I samme meddelelse sagde Google, at LiteRT, tidligere TensorFlow Lite, bliver en del af PyTorch-økosystemet: LiteRT-Torch-biblioteket konverterer en PyTorch-model med et enkelt litert-torch.convert() kald, og de konverterede modeller kører på Android, iOS, macOS, Windows, Linux, internettet og IoT-enheder. En pakke TPU-agenter, der skal automatisere ibrugtagning, diagnosticering og finjustering af inferensydelsen for PyTorch-modeller, er i en tidlig forhåndsvisning på GitHub.

Meddelelsen henviser udviklere til Google Clouds TPU-dokumentation, hvor de kan komme i gang med at køre PyTorch på TPU’er.