Ugrás a tartalomra
Neural Daily

EszközökHír

TorchTPU: natív PyTorch TPU-n, Hugging Face támogatással

A TorchTPU új PyTorch backend, amellyel a TPU ugyanolyan device lesz, mint a GPU. A Hugging Face bejelentette, hogy a Transformers, a Diffusers és a Kernels már natívan fut rajta, így egy Hub-modellt elég a "tpu" eszközre küldeni.

Írta: Mikola Máté Patrik · · 3 perc olvasás

Hugging Face TorchTPU banner: the Hugging Face emoji and the PyTorch logo as cartoon characters cheering behind a glowing Google TPU chip
Kép: Hugging Face · Sajtókép · Official Hugging Face blog image · Forrás

Mi változik?

Eddig aki PyTorch-kódot akart TPU-n futtatni, annak a PyTorch/XLA (torch_xla) maradt, vagy megtanulta a JAX-ot. A torch_xla működött, de lazy tenzorokkal dolgozott: a kód nem mindig akkor futott, amikor várnánk, apró változtatások is újrafordították a gráfot, a debug pedig nehézkes volt. A TorchTPU (GitHub) ezzel szemben a PyTorch natív backendje, a TPU egyszerűen egy újabb device lesz (Hugging Face).

A csomag TPU-ra hangolt Pallas kerneleket, elosztott tanítást és torch.compile integrációt hoz. TPU VM-en egy paranccsal telepíthető: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. Az extra index a CPU-s torch buildet hozza, így nem jön le feleslegesen a jóval nagyobb CUDA-s változat.

Mi fut már rajta?

  • Transformers v5: device_map="tpu" és a modell közvetlenül a TPU-ra töltődik. Alapból eager módban fut, torch.compile-lal gyorsítható, elosztott beállításban is használható.
  • Diffusers: a pipe.to("tpu") elég a prototípushoz, deployhoz egy compile hívás jön még. Nincs külön TPU fork. Az egy chipre túl nagy modelleket TensorParallelConfig osztja szét a chipek között betöltés közben.
  • Kernels: a Kernels Hubon kész TPU kernelek vannak, Pallasban írva, a Tokamax könyvtárból is behúzhatók.
  • Optimum TPU: a natív támogatás miatt Transformershez és Diffusershez már nem kell.

A számok

A Hugging Face mérése szerint egy Tokamax ragged-dot kernel MoE modellen nagy batch-nél sokat hoz: 64-es batch mellett a Qwen3-30B-A3B dekódolása 21,5-szer, a Gemma 4 26B-A4B-é 12,8-szor gyorsabb az alapértelmezett implementációnál, TPU v6e-8-on.

Hugging Face chart: decode throughput on TPU v6e-8, the Tokamax ragged_dot Hub kernel is 21.5x faster on Qwen3-30B-A3B and 12.8x faster on Gemma 4 26B-A4B at batch 64
Hugging Face · Sajtókép · Official Hugging Face blog image · Forrás

Forrás: Hugging Face. Dekódolási throughput TPU v6e-8-on (TP=8), Hub kernel és az alapértelmezett grouped_mm összevetése.

Diffusersben, Cloud TPU v6e-4-en, négy chipre osztva a Qwen-Image-2.1 egy 1024×1024-es, 40 lépéses képet compiled módban 3,5, eager módban 8,0 másodperc alatt generál. A MiniMax-H3 egy 5,2 másodperces, hangos videót 188, illetve 262 másodperc alatt készít el.

Mennyire stabil?

Még nem minden zöld. A TPU-n futó Transformers tesztek 96,5%-a megy át. A Hugging Face eddig 34 hibát jelentett a TorchTPU-nak, ebből 19 már javítva, a Transformersben is találtak hibákat (StaticCache, Qwen3 tensor-parallel tervek), ezeket upstream javították. Nightly TPU CI is készül, az eredmények a Transformers CI dashboardra kerülnek.

Mit jelent ez a fejlesztőknek?

A TPU eddig kerülőút volt PyTorch-osoknak, most kb. egy device-csere. Aki Google Cloudon tanít vagy nagy skálán szolgál ki modellt, annak érdemes összevetnie a TPU-t a GPU-s ajánlatokkal, mert a kódot már nem kell hozzá portolni. A Diffusers részletes receptje a TPU guide-ban van. A 96,5% viszont azt is jelenti, hogy élesítés előtt a saját modellen érdemes végigfuttatni a teszteket.

    MI segítségével készült, szerkesztő ellenőrizte. Hogyan készülnek a cikkeink

    Szerkesztőségünk ellenőrizte és szerkesztette. Rólunk

    Kapcsolódó cikkek