TorchTPU: könnyebb út a PyTorch-kódtól a Google TPU-kig
A Google TPU-jára költöző fejlesztőknek eddig gyakran más programozási logikához kellett alkalmazkodniuk. A TorchTPU ezt próbálja egyszerűbbé tenni: a Google gyorsítóchipjeit a PyTorch megszokott eszközei közé illeszti, a CPU és a CUDA-t használó GPU-k mellé.
A TorchTPU tehát nem új chip, hanem egy nyílt forrású PyTorch-eszközháttér. A Google és a Meta PyTorch-csapata közösen fejlesztette, nyilvánosan 2026. október 20-án jelent meg. A célja az, hogy a fejlesztők kisebb átírással futtathassanak PyTorch-modelleket Cloud TPU-kon.
Sok kód elindulhat, a gyorsaság viszont külön feladat
A legfontosabb ígéret egyszerű: sok esetben a model.to("cuda") helyett elég lehet a model.to("tpu") használata. A kód így ugyanazt a megszokott PyTorch-mintát követheti, amelyet CPU-n vagy GPU-n is használna. A program interaktívan futtatható és hibakereshető, a teljesítmény szempontjából fontos részeket pedig a torch.compile fordíthatja TPU-ra.
Ez jelentős könnyítés a korábbi PyTorch/XLA megoldáshoz képest. Annál a fejlesztőknek a késleltetett végrehajtás, a gráfépítés és az explicit szinkronizáció sajátosságaival is számolniuk kellett. A PyTorch/XLA dokumentációja szerint a TorchTPU váltja le ezt a projektet, amelynek archíválását 2026. november 15-re tervezik.
A könnyebb indulás ugyanakkor nem jelent automatikus sebességet. Egy GPU-ra gondosan hangolt rendszerből nem lesz pusztán eszköznév-cserével gyors TPU-s rendszer. A modell alakzatait, a párhuzamos feldolgozást és egyes műveleteket a TPU sajátosságaihoz kell igazítani. A saját CUDA-kódot, CUDA-specifikus kiterjesztéseket vagy GPU-ra írt egyedi kerneleket pedig külön át kell dolgozni.
Erős műszaki alap, fiatal fejlesztői környezet
A TorchTPU saját műveletimplementációkat, fordítási utat, fordítási gyorsítótárat, valamint elosztott tanításhoz szükséges képességeket tartalmaz. A nyilvános tárban Llama-, Qwen-, ResNet-, DLRM-, minGPT- és UNet-példák is elérhetők.
A projekt első publikus kiadása a torch_tpu 0.2.0, és Python 3.11 vagy újabb verziót céloz. Ezért technológiailag ígéretes, nagy léptékben már bemutatott rendszernek tekinthető, de fejlesztői környezete még jóval kevésbé kiforrott, mint a CUDA-é. A hibajavítási rutinoknak, külső integrációknak és gyakorlati tapasztalatoknak még idő kell.
A forráskód Apache 2.0 licenc alatt érhető el, így ellenőrizhető és módosítható. A rendszer ugyanakkor továbbra is a Google infrastruktúrájához kötődik: célhardvere a Cloud TPU, a futtatás pedig a Google TPU-s szoftverrétegére támaszkodik.
A teljesítményt még óvatosan kell értékelni
A Google és a Meta saját mérése szerint a TorchTitan egy Llama 3 405B modell tanításánál 1024 TPU v7x chipen, 8K kontextushosszal és BF16 formátumban chpenként 500 token/másodperc feletti teljesítményt, valamint 50 százalék feletti modell-FLOPS-kihasználtságot ért el. Ez azt jelzi, hogy a rendszer képes nagyon nagy méretben működni.
Ezek azonban a projekt fejlesztőinek eredményei. Nem jelentenek általános CUDA-győzelmet, mert a valódi összevetéshez azonos modell, beállítás, késleltetési cél és teljes rendszerköltség kellene. A TPU különösen olyan nagyméretű, ismétlődő feladatoknál lehet erős, amelyek jól eloszthatók sok chip között. Kevésbé kedvező helyzetet teremthet a sok változó alakzat, az egyedi GPU-kód vagy a CUDA-specifikus komponensek nagy száma.
Kezd formát ölteni az ökoszisztéma
Az induláskor a TorchTPU már kapcsolódott több elterjedt eszközhöz: a Hugging Face Transformershez, Diffusershez és Kernelshez, továbbá a vLLM-hez, SGLanghez, TorchTitanhez, TorchRechez és Rayhez. A Transformers v5-ben például a modell device_map="tpu" paraméterrel tölthető TPU-ra.
Ezek az integrációk azt jelentik, hogy a már ismert modellkönyvtárakat és kiszolgálóeszközöket használó csapatok kisebb újratanulással próbálhatják ki a TPU-kat. A Meta PyTorch-csapata társfeljesztőként vesz részt a munkában, a Hugging Face pedig natív TPU-támogatással kapcsolódott be. A vLLM és az SGLang TPU-változatai a nagy volumenű modellkiszolgálást célozzák.
A TorchTPU fő előnye tehát az átállás egyszerűsítése, fő korlátja pedig a Google TPU-infrastruktúrához kötöttség és a még épülő fejlesztői környezet. Új, nagyléptékű PyTorch-projektnél érdemes számolni vele, különösen elosztott tanításnál és nagy nyelvi modellek kiszolgálásánál. Egy meglévő, erősen CUDA-specifikus rendszernél viszont először szűk körű tesztből derülhet ki, mennyibe kerül és mit hoz az átállás.