A Hugging Face natív PyTorch-támogatást hoz a TPU-kra

Egyszerűbb TPU-használat PyTorch mellett

A Hugging Face október 20-án jelentette be a TorchTPU-t. A vállalat szerint a megoldással a TPU a PyTorch számára ugyanolyan eszközzé válik, mint bármely más támogatott végrehajtási célpont. A felhasználó egy, a Hugging Face Hubról kiválasztott modellt a tpu eszközre küldhet.

A TPU, vagyis Tensor Processing Unit a Google saját mesterségesintelligencia-gyorsítója. Ezeket a chipeket elsősorban a neurális hálózatokban gyakori mátrixműveletekre tervezték. A Google a forrás szerint körülbelül tíz éve fejleszti őket, és többek között saját modelljeinek, köztük a Gemininek a tanítására és kiszolgálására használja. A TPU-k a Google Cloudon egyetlen chiptől több ezer chipből álló fürtökig érhetők el.

A TPU-k szoftveres környezete eddig elsősorban a TensorFlow és a JAX köré épült. A PyTorch használatához a fejlesztőknek jellemzően a PyTorch/XLA, vagyis a torch_xla köztes megoldásán kellett átmenniük. A Hugging Face szerint ez lusta tenzorokat, újrafordításokat és nehezebben követhető hibakeresést eredményezhetett.

A Transformers és a Diffusers is közvetlenül futtatható

A TorchTPU új hátteret ad a PyTorchhoz. A csomag TPU-kra hangolt Pallas kerneleket, elosztott tanítási támogatást és torch.compile integrációt kínál. TPU VM-en a telepítéshez a Hugging Face egyetlen parancsot javasol: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. A vállalat szerint az extra csomagforrás azért szükséges, mert így a TPU-hoz nem szükséges, nagyobb CUDA-verzió helyett a CPU-s PyTorch-csomag települ.

A Transformers v5 modellbetöltésénél a device_map="tpu" beállítás közvetlenül a TPU-ra helyezi a modellt. A futtatás alapértelmezés szerint azonnali módban történik, de a fejlesztők a jobb teljesítményért fordítást, illetve elosztott futtatást is választhatnak.

A Diffusersben a TPU első osztályú eszközként jelenik meg. Prototípus-készítéshez és hibakereséshez elegendő a csővezetéket a pipe.to("tpu") hívással áthelyezni. Telepítésnél a fordítás egy további hívással kapcsolható be. A Hugging Face szerint a meglévő csővezetékek és kódok változatlanul használhatók, így nincs szükség külön TPU-változat fenntartására.

A több chipet igénylő modellekhez a TensorParallelConfig használható. Ezzel a transzformátor betöltés közben több chip között osztható fel, így a teljes modellnek nem kell egyetlen chipen elférnie.

Kernelek, mérések és folyamatos tesztelés

A Kernels Hub kész TPU-kerneleket kínál. Ezek Pallasban készülnek, amely alacsony szintű hozzáférést ad a hardverhez, és a Hugging Face szerint teljesítménynövekedést hozhatnak. Tokamaxból származó kernelek is használhatók, többek között kevert szakértői modellekben.

A Hugging Face mérései szerint 64-es kötegméretnél a Hub kernelje a Qwen3-30B-A3B modell dekódolását 21,5-szeresére, a Gemma 4 26B-A4B modellét pedig 12,8-szeresére gyorsította a TPU v6e-8-on, az alapértelmezett megoldáshoz képest.

A fejlesztők a Transformers tesztcsomagjával folyamatosan ellenőrzik a TPU-támogatást. A futtatható TPU-tesztek 96,5 százaléka teljesül. Eddig 34 hibát jelentettek a TorchTPU tárházában, ezek közül 19-et már javítottak. A tesztelés a Transformersben is talált hibákat, többek között a StaticCache és a Qwen3 tenzoros párhuzamosítási tervei körül.

A bejelentés a PyTorchot használó Hugging Face-felhasználóknak egyszerűbb belépést kínál a Google Cloud TPU-k világába. A TorchTPU telepítése után a modellek és könyvtárak natív TPU-futtatása a forrás szerint közvetlenül elérhetővé válik.