Google heeft TorchTPU, de open source-stapel die het mogelijk maakt om PyTorch native uit te voeren op hun TPU’s, vrijgegeven. Deze AI-versnellers worden gehuurd via Google Cloud. De aankondiging deed Google op 20/10/2026 tijdens de PyTorch Conference 2026. Het project, ontwikkeld in samenwerking met het PyTorch-team van Meta, staat op GitHub onder de Apache 2.0-licentie. Google beweert dat met TorchTPU TPU’s nu een native hardware-optie voor PyTorch worden, net als GPU’s en CPU’s.
De kernpunten van TorchTPU in 20 seconden
- Google maakt TorchTPU met een Apache 2.0-licentie open, samen met het PyTorch-team van Meta ontwikkeld.
- Voert PyTorch uit op TPU met minimale codewijzigingen, zowel in eager-modus als met torch.compile.
- vLLM, SGLang, Hugging Face, Ray en llm-d integreren het al.
- Google spreekt over bijna lineaire schaalbaarheid tot 8.192 TPU Ironwood.
De aankondiging, ondertekend door Mark Lohmeyer, vice-president en algemeen manager van AI-infrastructuur en computing bij Google Cloud, en Bill Jia, vicepresident engineering, richt zich op een van de klassieke obstakels bij het gebruik van TPU: de meeste AI-werkbelastingen draaien op PyTorch, terwijl de software van Google voor hun chips traditioneel gebaseerd is op JAX en TensorFlow. Voor cloud-gebruikers is de vraag of ze workloads kunnen verplaatsen naar een andere versneller zonder code te herschrijven.
Wat verandert er ten opzichte van PyTorch/XLA?
Tot nu toe was de enige manier om PyTorch op TPU te gebruiken via PyTorch/XLA, beschreven door het Google-ontwikkelaarsblog als het voorganger van TorchTPU. XLA ondersteunde alleen pure SPMD-code (hetzelfde programma op alle knooppunten). TorchTPU integreert in PyTorch via de PrivateUse1-interface, die het framework reserveert voor externe hardware backends. Hierdoor zijn tensors op TPU gewone PyTorch-tensors, zonder wrappers of subklassen.
Google stelt voor om in eager-modus te debuggen, operationeel per operatie, en daarna over te schakelen naar torch.compile voor productie. Bij die tweede stap vangt TorchTPU de graf op met TorchDynamo en compileert deze met XLA in plaats van TorchInductor, de standaard compiler van PyTorch. Volgens Google is er daarnaast een eager-stand met automatische fusie van operaties dat tussen de 50% en meer dan 100% verbetering biedt ten opzichte van standaard eager. Voor gedistribueerde training ondersteunt het DDP, FSDPv2 en DTensor, en kan het kernels gebruiken geschreven in Pallas of JAX.
De versie op PyPI is 0.2.0, vereist Python 3.11 of hoger en installeert via pip.
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpuSchaalbaarheid op Ironwood en inferentie-stack op GKE
De prestatiecijfers afkomstig van Google zijn gebaseerd op eigen tests. Met TorchTitan, de referentie-implementatie van PyTorch voor gedistribueerde training, beweert Google 52% benutting van de FLOPs (MFU) te halen bij het trainen van Qwen 3 8B, met bijna lineaire en stabiele schaalbaarheid tot 8.192 TPU Ironwood (TPU v7x), de serie die Google introduceerde voor grootschalige inferentie. TorchTitan ondersteunt onder andere modellen als Llama 3 met 70.000 en 405.000 miljard parameters, Qwen 3 235B met MoE-architectuur en FLUX.1 12B.
Voor inferentie heeft Google native backends ontwikkeld voor vLLM, samen met Inferact, en voor SGLang, samen met RadixArk. Daarnaast ondersteunt het llm-d 1.0, een project in de Sandbox-fase van de CNCF waarin Red Hat, IBM en meer dan 500 collaborateurs van 25 organisaties deelnemen. Dit project voegt dynamisch cache-georiënteerd routing en scheiding van prefill- en decode-fasen toe. Op GKE, de Kubernetes-dienst van Google Cloud, kunnen modellen worden bediend op GPU of TPU via hetzelfde endpoint.
Als voorbeeld citeert Google Zenteiq.ai, waarvan de oprichter zegt dat ze de totale inferentieprestatie met GKE Inference Gateway en llm-d meer dan vier keer hebben verhoogd, van 1.470 naar 5.970 miljoen tokens per dag per 1.000 TPU v6e versnellers. Ray, de orkestrator van Anyscale, ondersteunt inmiddels ook TorchTPU voor workload-distributie tussen CPU, GPU en TPU.
Een alternatief voor CUDA, met nuances
Deze beweging vertegenwoordigt een duidelijke verandering op de acceleratormarkt. Veel van NVIDIA’s voordeel zit in CUDA en de optimalisatie van PyTorch op hun GPU’s. TPU’s concurreren nu op prijs en prestaties, zoals besproken bij de analyse van Google TPU’s versus NVIDIA. Echter, software bleef een barrière. Meta, dat TPU’s van Google gebruikt, is geïnteresseerd in minder afhankelijkheid van één leverancier. PyTorch-engineer Jana van Greunen verklaarde dat TorchTPU hen in staat stelt om PyTorch-workloads op TPU’s te draaien voor taalmodellen, advertentie-systemen en recommendation engines, zonder vast te zitten aan één hardware-architectuur.
Er zijn nog zaken die buiten Google’s tests moeten worden aangetoond. Het ontwikkelplan omvat onder andere het verminderen van hercompilaties bij veranderingen in sequentielengte of batchgrootte, volledige ondersteuning voor dynamische vormen in torch.compile en het uitbreiden van de bibliotheek met vooraf gecompileerde kernels. Dit zijn bekende problemen voor XLA-gebaseerde compilers, vooral bij variabele inferentielasten. Bovendien werkt TorchTPU alleen met Google TPU’s die via Google Cloud gehuurd worden. De portabiliteit is dus beperkt tot tussen GPU en TPU binnen PyTorch en niet tussen verschillende cloudaanbieders.
Veelgestelde vragen
Wat is TorchTPU?
De softwarestack van Google om PyTorch native op TPU’s uit te voeren. Het is open source, onder de Apache 2.0-licentie, en samen ontwikkeld met het PyTorch-team van Meta.
Vervangt TorchTPU PyTorch/XLA?
Google beschouwt PyTorch/XLA als haar voorganger. TorchTPU integreert als een native backend in PyTorch, maar Google heeft niet officieel aangekondigd dat PyTorch/XLA wordt stopgezet.
Wat voor prestaties biedt TorchTPU?
Volgens Google haalt het 52% MFU bij het trainen van Qwen 3 8B met TorchTitan en bereikt het bijna lineaire schaalbaarheid tot 8.192 TPU Ironwood. Dit zijn de eigen cijfers van Google.
Kan TorchTPU buiten Google Cloud worden gebruikt?
De software is open source, maar werkt alleen op TPU’s die via Google Cloud zijn gehuurd.
Bronnen: