A TorchTPU bevezetésével a TPU immár natív eszközként használható PyTorch alatt, így a korábbi torch_xla-hoz való portolás vagy a JAX megtanulása már nem kötelező. Ez lehetővé teszi, hogy a Hugging Face ökoszisztéma komponensei — például a Transformers, Diffusers és Kernels — közvetlenül TPUn fussanak: válassz egy modellt a Hubról, küldd "tpu" eszközre, és a kódod elindul.

Mi az a TPU?

A TPU (Tensor Processing Unit) a Google saját AI gyorsítója, kifejezetten a neurális hálózatok mátrixszorzásaira optimalizált chipek. A Google több mint tíz éve fejleszti ezeket, és saját nagy modelljeinek képzésére és üzemeltetésére is használja, például a Gemini modellhez. A TPU-k a Google Cloudon elérhetők, egyetlen chipestől egészen több ezer chipet tartalmazó podokig. Különösen akkor erősek, ha nagy számítási kapacitásra van szükség tréninghez vagy nagy skálájú kiszolgáláshoz.

A nehézség eddig a szoftveres támogatás körül adódott: a TPU stack elsősorban TensorFlow és JAX környezetre épült, ezért a PyTorch-felhasználóknak többet kellett bajlódniuk a működésért.

Az "öreg" út: torch_xla

Korábban a PyTorch TPUn való futtatása jellemzően a PyTorch/XLA (torch_xla) használatát jelentette. Ez működött, de sajátosságai voltak: lusta (lazy) tenzorok miatt a kód nem mindig futott akkor, amikor a fejlesztő azt várta; kis változtatások újrafordításokat válthattak ki, és a hibakeresés nehezebb volt.

A Hugging Face korábban az Optimum TPU csomagot építette ki, hogy modellek inference- és tréningtámogatást kapjanak TPUn, de most, a TorchTPU megjelenésével az Optimum TPU már nem szükséges a Transformers és Diffusers modellek futtatásához TPUn.

Mit hoz a TorchTPU?

A TorchTPU új backendként adja hozzá a TPU-t a PyTorchhoz. A fejlesztő hagyományos PyTorch-kódot írhat, miközben a háttérben a TorchTPU Pallas kernelokat biztosít a TPU-khoz optimalizált műveletekhez. Támogatja a disztribúciós tréninget, és integrálható a torch.compile mechanizmussal, ha nagyobb sebességet szeretnénk.

A telepítés egy TPU VM-en egyszerű parancs:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

Az extra index a torch CPU buildjét adja, enélkül pip a nagyobb CUDA buildet töltené le, amelyre TPUn nincs szükség.

Hugging Face könyvtárak TPUn

Transformers

A Transformers v5 alapgondolata, hogy a modellik definiálása legyen az egyetlen hiteles forrás, így a modellek bárhol futhassanak. A "bárhol" mostantól a TPU-t is magában foglalja. Egy modell betöltéséhez TPUn elegendő a device_map="tpu" megadása; a modell ezzel közvetlenül a TPU-n lesz betöltve. A modell alapértelmezetten eager módban fut; opcionálisan használható torch.compile nagyobb teljesítményért, illetve disztribúcióban is futtatható.

Diffusers

A Diffusers könyvtárban a TPU immár elsőrendű eszköz. Prototípus-készítéshez vagy hibakereséshez az eager mód használatához elegendő a pipe.to("tpu"). Telepítésre, éles futtatásra a compiled mód egy compile hívással gyorsítható. A pipeline-ok és a kód megmaradnak változatlanul, nincs szükség külön TPU-forkra.

Példakódrészlet a StableDiffusion3Pipeline használatára, ahol a csővezeték dtype-ja torch.bfloat16 és TPUn fut:

import torch from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16) pipe.to("tpu")

eager mód: ez a teljes változtatás

opcionális: compiled mód

pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

Az első hívás lefordítja a grafot, és ugyanazzal a képmérettel minden további hívás újrafelhasználja a lefordított grafot.

Nagy modellek több chipen

Ha egy modell túl nagy egyetlen chiphez, a TensorParallelConfig átadásával a from_pretrained során a transzformert a betöltés közben szét lehet shardolni a chippek között, így a teljes modell soha nem kell, hogy egy helyen elférjen. Sok pipeline alapból támogatja ezt, és ahol nincs, ott egy tp_plan hozzáadása megoldást jelent. A Diffusers TPU útmutató részletezi a szükséges lépéseket.

Példák Cloud TPU v6e-4 kluszteren futtatott eredményekről (TP=4, bf16, batch 1):

• MiniMax-H3 (video + audio) és Qwen-Image-2.1 példák különböző képméretekkel és lépésszámokkal; egyes esetekben a compiled futtatás lényegesen gyorsabb, máskor az eager mód is használható. (A cikk konkrét időadatokat és összehasonlítást közöl egyes promptokra: például 1344×768, 5.2 s, 49 lépés; 1024×1024, 40 lépés: 3.5 s compiled vs. 8.0 s eager; stb.)

Kernels

A Kernels Hub kész TPU kernelokat tartalmaz, amelyeket a Kernels könyvtárral tölthetsz be és használhatsz a modell gyorsítására. A TPU kernelok Pallas nyelven íródtak, amely alacsony szintű vezérlést ad a hardver felett. Emellett importálhatóak Tokamax könyvtárból származó kernelok is.

A cikk egy ragged-dot kernelt mutat be Tokamax-ból, amelyet egy Mixture-of-Experts modellbe integrálnak; a példában a kernelt regisztrálják és az experts_forward funkcióként használják egy Qwen nagy modellnél.

Teljesítmény: a nagyobb batch-eknél nagyobb a nyereség. Például batch méret 64 esetén dekódolásban 21.5× gyorsulás volt megfigyelhető a Qwen3-30B-A3B modellen és 12.8× Gemma 4 26B-A4B-n a hub kernel vs. alapértelmezett implementáció összehasonlításában. A cikk említi a dekódolási áteresztőképesség méréseit TPU v6e-8-on hub kernel és alapértelmezett implementáció között.

Tesztelés és folyamatos fejlesztés

A Hugging Face folyamatosan fejleszti a TPU-támogatást a transformers könyvtár tesztkészletének futtatásával. Dolgoznak egy éjszakai (nightly) CI beállításon TPUn, és terveik szerint a futtatási eredmények megjelennek majd a Transformers CI dashboardon. Jelenleg a TPUn futó tesztek 96.5%-a sikeres. Minden talált hibáról készítenek kis reprodukáló szkriptet és issue-t nyitnak a TorchTPU repozitóriumban: eddig 34 issue-t jelentettek, amelyek közül 19 javítva is lett. A TPU-s futtatás egyébként hibákat is talált a Transformers saját kódjában (például a StaticCache és a Qwen3 tensor-parallel tervek esetén), amelyeket szintén javítottak upstream.

Hogyan kezdj hozzá

A cikk konklúziója szerint a TPU-k immár nem mellékes opciók a PyTorch-felhasználóknak. A lépések rövidek: telepítsd a TorchTPU-t, válassz egy modellt a Hubról, és küldd device_map vagy .to("tpu") segítségével a TPU-ra. A TorchTPU GitHub-on, a Kernels Hub-on és a Cloud TPU dokumentációban találsz további forrásokat és útmutatókat. Ha problémába ütközöl, a szerzők azt javasolják, hogy nyiss issue-t és oszd meg a tapasztalataidat.

A Hugging Face csapata bíztatja a felhasználókat, hogy próbálják ki az új támogatást, és várják a visszajelzéseket.