Google har gjort TorchTPU tillgängligt som öppen källkod. Det är ett mjukvarulager som gör det möjligt för utvecklare att köra PyTorch – det brett använda verktyget för att bygga och köra AI-modeller – på Googles Tensor Processing Units, eller TPU:er. TorchTPU har utvecklats tillsammans med Metas PyTorch-team och är utformat för att befintliga arbetsbelastningar ska kunna flyttas till TPU-hårdvara med relativt små kodändringar, i stället för att skrivas om från grunden.
Enligt Google stöder TorchTPU både eager-exekvering, där operationer körs efter hand som de anropas och som är användbart vid felsökning, och kompilerad exekvering för produktionsarbetsbelastningar. Utvecklare kan ändra en enhetsinställning från ”cuda” till ”tpu” och använda TPU:ns distribuerade backend. Google hänvisar också till PyTorches API för att identifiera acceleratorer som ett alternativ. TorchTPU finns på GitHub, men på utvecklingsplanen återstår bland annat arbete med dynamiska indatastorlekar, kompileringshastighet, stöd för fler operatorer, förstärkningsinlärning och kvantisering.
Vad innebär det för utvecklare?
Lanseringen kopplar TPU:er till programvara som många PyTorch-utvecklare redan använder. Google uppger att TorchTPU fungerar med serveringsverktygen vLLM och SGLang samt med projekt för träning och orkestrering, däribland TorchTitan, Ray och TorchRec. I Transformers v5 säger Hugging Face att utvecklare kan välja en TPU med inställningen device_map="tpu". Även biblioteket Diffusers stöder körning på TPU, inklusive kompilerat läge och fördelning av arbetet över flera acceleratorer.
Det här är en annan lösning än PyTorch/XLA, Googles tidigare sätt att köra PyTorch på TPU:er. Enligt Google krävde det systemet mer omfattande kodändringar, hade långa kompileringstider och saknade stöd för PyTorches eager-läge. Google har byggt en testpipeline som enligt InfoQ utvärderar omkring 5 000 Hugging Face-modeller inom fler än 500 arkitekturer. Det tyder på bred testning, men innebär inte att alla modeller eller arbetsbelastningar garanterat fungerar utan ändringar.
Hur står det sig mot andra hårdvarualternativ?
Google uppger att TorchTPU:s DeferAndFuse-läge – en kombination av eager- och kompilerad exekvering – kan vara minst 50 % snabbare än Strict Eager-läget och samtidigt minska den maximala minnesanvändningen hos acceleratorn. I ett separat test körde Google en Muse Glimmer-modell med 30 miljarder parametrar i standardversionen av PyTorch. InfoQ rapporterar en inferenshastighet på 25 700 token per sekund på ett TPUv7-chip vid en sekvenslängd på 5 120, och 22 400 vid 10 240. Träning på fyra TPUv7-chip nådde 5 600 respektive 4 600 token per sekund vid dessa sekvenslängder. Det här är resultat som Google har rapporterat, inte en oberoende jämförelse.
TechCrunch jämför TorchTPU med Amazons TorchNeuron, som presenterades samma vecka och ska användas för att köra PyTorch-arbetsbelastningar på Amazons Trainium-chip. Båda satsningarna ska göra det enklare att använda alternativ till Nvidias GPU:er. TorchTPU ger PyTorch-användare ytterligare ett hårdvarualternativ, men lanseringen visar inte i sig att en flytt av en arbetsbelastning sänker kostnaderna eller ger bättre prestanda än utvecklarens nuvarande lösning.
Varför det spelar roll
För utvecklare som redan använder PyTorch kan TorchTPU göra det enklare att testa eller driftsätta arbetsbelastningar på Googles TPU:er utan att bygga om dem från grunden. Det är fortfarande oklart om en sådan flytt förbättrar kostnaden eller hastigheten för en viss modell; de prestandasiffror som anges kommer från Google självt.
