Google kondigde TorchTPU op 20 oktober 2026 aan als open-sourcebackend waarmee PyTorch-code op Google TPU’s kan draaien. Voor een basisrun zijn installatie en een apparaatkeuze nodig; gedistribueerde uitvoering vraagt daarnaast een TPU-backend en specifieke startinstellingen. Meta werkte mee aan de ontwikkeling.
Amazon kondigde een dag eerder TorchNeuron voor zijn Trainium-accelerator aan, meldde TechCrunch op 20 oktober 2026. TorchTPU richt zich op PyTorch-workloads op Google TPU’s.
Een basisrun begint met installatie en apparaatkeuze
Google documenteert de installatie met pip install torch_tpu en de keuze van het apparaat via torch.device("tpu"). In een voorbeeld van het PyTorch-project verandert .to("cuda") in .to("tpu"). Eager execution, waarbij PyTorch bewerkingen direct uitvoert, is de standaard. Voor gecompileerde uitvoering is torch.compile beschikbaar.
Gedistribueerde uitvoering vraagt TPU-specifieke instellingen
Bij een gedistribueerde workload gebruikt TorchTPU de backend tpu_dist en één proces per apparaat. De startconfiguratie moet ook informatie over de TPU-topologie en de slice builder bevatten, of via de gedocumenteerde helper worden ingesteld.
| Workload | Apparaat of backend | Uitvoering en startinstellingen |
| Basisrun | Installeer torch_tpu en kies torch.device("tpu") of .to("tpu") | Eager execution is standaard; torch.compile is beschikbaar voor gecompileerde uitvoering |
| Gedistribueerde workload | Backend tpu_dist, met één proces per apparaat | Stel TPU-topologie en slice-builderinformatie in, of gebruik de gedocumenteerde helper |
Dat verschil is relevant zodra een toepassing over meerdere TPU’s wordt verdeeld: dan volstaat alleen het wijzigen van de apparaatkeuze niet. Ook voor een eenvoudige run garandeert die wijziging geen topprestaties. Volgens het PyTorch-project kunnen onder meer tensorvormen, sharding en bewerkingen afstemming op de TPU vereisen.
Wat de gepubliceerde tests zeggen over compatibiliteit
Het PyTorch-project rapporteerde tests op geselecteerde modellen. In Tier 1 draaiden 4.821 van de 4.917 modellen zonder crash, oftewel 98%. In Tier 2 bereikten 638 van de 657 representatieve modellen numerieke pariteit met een GPU-baseline, oftewel 97%. Bij een test op begrensde dynamiek waren 834 van de 839 geëvalueerde causale LLM’s correct, gerapporteerd als 99%.
Deze testgroepen geven geen garantie dat iedere PyTorch-workload werkt. Het project noemt ongeveer 450 ondersteunde ATen-operators; structurele of samengestelde delegates tellen niet mee in dat aantal. In één GPU/CUDA-afkomstige Muse Glimmer 30B-implementatie werd FlexAttention, dat TorchTPU in die test niet ondersteunde, vervangen door F.scaled_dot_product_attention.
Het project rapporteerde ook een trainingstest met Qwen3 MoE 235B op maximaal 8.192 TPU v7-chips. Bij die configuratie, met een batchgrootte van 262.144 en een sequentielengte van 4.096, meldde het 2.104 tokens per seconde per TensorCore, 27,1% MFU en 31,2 seconden per stap. Dit is een projectrapport over één specifieke configuratie, geen gestandaardiseerde vergelijking met een GPU voor uiteenlopende workloads.
Bestaande modelgewichten laden op een TPU
Als de modeldefinitie en de gebruikte bewerkingen compatibel zijn, kunnen opgeslagen modelgewichten via een state_dict in het model worden geladen. PyTorch’s torch.load ondersteunt het omzetten van de apparaatlocatie van opgeslagen tensors; daarna kan het model de TPU als apparaat gebruiken. Dat biedt een route om met geschikte gewichten op een TPU verder te werken.
