- Det är en inbyggd PyTorch-backend som låter dig köra kod direkt på Cloud TPU:er utan att skriva om den eller använda JAX.
- Den introducerar tre förkörningslägen (Strict, Debug och DeferAndFuse) och optimering via torch.compile.
- Den integreras med ledande verktyg som vLLM, SGLang, Hugging Face och Ray för att skala träning och inferens.
- Det syftar till att minska beroendet av NVIDIAs CUDA-arkitektur genom att underlätta migreringen av arbetsbelastningar till TPU-hårdvara.
Om du arbetar inom artificiell intelligens vet du att det länge fanns en sorts osynlig vägg som skilde PyTorch-ekosystemet från Googles hårdvara. Fram tills nyligen innebar det att man antingen helt skulle se över sin kod eller byta till JAX för att kunna utnyttja kraften hos TPU:er, något som inget utvecklingsteam är förtjust i när de väl har sina modeller igång. TorchTPU har kommit för att bryta ner den väggen , vilket gör det möjligt för utvecklare att driftsätta sina PyTorch-projekt på Googles Tensor Processing Units utan att behöva skriva om en enda rad med kärnlogik.
Den här utgåvan är inte bara en enkel uppdatering, utan ett paradigmskifte som utvecklats tillsammans med PyTorch-teamet på Meta. I grund och botten är det en nativ backend som ersätter den gamla PyTorch/XLA , vilket eliminerar behovet av att bränna lata grafer och äntligen låter TPU:er kännas som förstklassiga medborgare tillsammans med NVIDIA GPU:er. Detta är ett mästerverk i att demokratisera tillgången till specialiserad hårdvara som, om den används korrekt, kan minska kostnaderna och dramatiskt öka effektiviteten i träning och härledning av massiva modeller.
Vad exakt är TorchTPU och hur fungerar det?
I grund och botten är TorchTPU en öppen källkods-programvarustack som gör det möjligt för PyTorch att kommunicera direkt med kiseln i Googles TPU:er. Dess kraftfullaste funktion är dess "Eager First" -filosofi , vilket innebär att kod körs omedelbart, precis som vi är vana vid i PyTorch, vilket avsevärt förenklar felsökning och interaktiv testning. För att uppnå detta är det integrerat via PrivateUse1-gränssnittet, vilket möjliggör hantering av vanliga PyTorch-tensorer utan att förlita sig på en JAX-körningstid.
För att säkerställa ett smidigt arbetsflöde erbjuder TorchTPU tre förkörningslägen:
- Strikt läge: Det är standardalternativet; det startar operationer asynkront enligt PyTorchs imperativa semantik.
- Felsökningsläge: Perfekt när saker går fel, eftersom den kör allt synkront och kontrollerar tensorgränser för att upptäcka minnesfel eller NaN.
- DeferAndFuse-läge: Det är här magin börjar, när den sammanfogar sekvenser av operationer till minimera HBM-minnestrafik, vilket uppnår prestandaförbättringar som kan variera från 50 % till över 100 % jämfört med strikt läge.
När utvecklaren behöver maximal prestanda för produktion kommer torch.compile med TPU-backend in i bilden . Denna process fångar beräkningsgrafen med hjälp av TorchDynamo och översätter den till StableHLO och sedan till XLA, vilket genererar optimerade binärfiler som lagras i en delad cache för att undvika irriterande omkompileringar.
Ett robust ekosystem för träning och inferens
Google har inte lanserat verktyget ensamt; istället har de omgivit det med en arsenal av integrationer som gör det användbart i verkliga miljöer. För de som arbetar med öppna viktningsmodeller är partnerskapet med Kramande ansikte är nyckelnNu behöver du bara lägga device_map="tpu" i Transformers v5 så att modellen körs på Googles hårdvara. Dessutom har CI/CD-pipelines skapats som testar Transformers-sviten varje natt för att säkerställa att allt går smidigt.
Inom inferensområdet utmärker sig två grundläggande motorer: vLLM-TorchTPU och SGLang-TorchTPU . Dessa backend-system möjliggör exekvering av modeller som Gemma 4, Qwen eller DeepSeek genom att utnyttja anpassade Pallas-kärnor. De stöder avancerade tekniker som spekulativ avkodning och kvantisering i FP8 eller MXFP4 , vilket möjliggör otroligt snabb tokengenerering utan att offra noggrannhet.
För molnorkestrering finns llm-d , ett CNCF-projekt som fungerar som hjärnan som distribuerar förfrågningar över Kubernetes. Dess prefix-cache-medvetna routningsfunktion är en pärla, eftersom den skickar varje förfrågan till den replik som redan har kontexten laddad, vilket undviker redundanta beräkningar och optimerar minnesanvändningen.
Massiv skalbarhet och banbrytande verktyg
När det gäller att träna jättemodeller är effektivitet allt. TorchTitan på TPU är riktmärkesplattformen för detta, och integrerar SimpleFSDP och optimeringar på grafnivå. Den har validerats för att skala nästan linjärt upp till 8 192 Ironwood TPU-chip , vilket uppnår en modell-FLOP-utnyttjandegrad på upp till 52 % på arkitekturer som Qwen 3 8B.
För de som behöver hantera komplexa arbetsflöden är integration med Ray avgörande . Ray hanterar TPU-sliceprovisionering, topologi och automatisk redundans, vilket gör att dataförbehandling och träning kan ske i ett enda, koordinerat flöde. Samtidigt optimerar TorchRec rekommendationssystem genom att delegera sökningar efter glesa inbäddningstabeller till TPU:ernas SparseCores, medan TensorCores hanterar den täta matematiken.
De har till och med övervägt utvecklare som inte är TPU-hårdvaruexperter med TPU Agent Suite . Dessa AI-agenter hjälper till att automatisera modellintroduktion, diagnostisera flaskhalsar genom att analysera XLA-spår och skriva optimerade Pallas-kärnor, vilket fungerar som en brygga till teknisk kunskap.
Påverkan på industrin och kampen om kisel
Historiskt sett har NVIDIA dominerat branschen tack vare CUDA, vilket skapat ett ekosystem där mjukvara och hårdvara var sammanflätade. TorchTPU försöker bryta det monopolet genom att eliminera migrationsfriktion . Om kostnaden för att byta leverantör inte längre handlar om månaders ingenjörskonst, utan bara en månadsfaktura, har företag ett verkligt incitament att diversifiera sin hårdvara.
Även om Meta fortfarande använder GPU:er, visar det faktum att de redan kör produktionsarbetsbelastningar med SGLang-TorchTPU teknikens mognad. Konkurrensen fokuserar nu på kostnaden per miljon tokens . Med ankomsten av den sjunde generationen av TPU:er (Ironwood) strävar Google efter att leverera energieffektivitet och prestanda som kommer att tvinga AI-jättar att se bortom traditionella GPU:er.
Denna nya arkitektur gör träning och driftsättning av modeller med miljarder parametrar mer lättillgängliga. Genom att integrera disaggregerade minnesfunktioner och optimerad överföring via TPU Sync har syntetisk datagenerering och multimodal inferens uppnått rekordbrytande token-per-sekund-prestanda, vilket positionerar TorchTPU som det mest seriösa och gångbara alternativet för dem som vill undvika absolut beroende av en enda chipleverantör.
Passionerad författare om bytesvärlden och tekniken i allmänhet. Jag älskar att dela med mig av min kunskap genom att skriva, och det är vad jag kommer att göra i den här bloggen, visa dig alla de mest intressanta sakerna om prylar, mjukvara, hårdvara, tekniska trender och mer. Mitt mål är att hjälpa dig att navigera i den digitala världen på ett enkelt och underhållande sätt.