Hugging Face oznámil nativní podporu TPU v knihovnách Transformers, Diffusers a Kernels prostřednictvím backendu TorchTPU. Vývojáři tak mohou spouštět modely z Hubu přímo z PyTorche. Pro práci s Transformers a Diffusers už nepotřebují Optimum TPU, který dosud zajišťoval podporu těchto modelů na TPU.
Běžný kód v PyTorchi místo dosavadní cesty přes XLA
Dosavadní běh PyTorche na TPU vedl přes PyTorch/XLA. Ten používal odložené vykonávání operací, takže se operace nemusely provést ve chvíli, kdy je kód vyvolal. Podle Hugging Face mohly i drobné změny vyvolat novou kompilaci výpočetního grafu a tento způsob práce komplikoval ladění.
TorchTPU přidává TPU jako další backend PyTorche a umožňuje používat běžný kód této knihovny. Součástí jsou jádra Pallas optimalizovaná pro TPU, podpora distribuovaného trénování a integrace torch.compile.
Transformers v5 načítá model přímo na TPU
V Transformers v5 se cílové zařízení určuje při načítání modelu parametrem device_map="tpu". Výchozím způsobem vykonávání je režim eager. Vedle něj lze použít torch.compile nebo model spustit v distribuovaném uspořádání; volba kompilace tak není podmínkou samotného načtení a spuštění na TPU.
Diffusers podporuje kompilaci i rozdělení mezi čipy
V Diffusers se pipeline pro prototypování či ladění v režimu eager přesune na TPU voláním pipe.to("tpu"). Pro nasazení v kompilovaném režimu přibývá jedno volání kompilace, zatímco kód pipeline zůstává stejný. Není nutné udržovat samostatnou větev určenou pouze pro TPU.
Kompilovaná pipeline vytváří výpočetní graf při prvním volání. Při dalších voláních ho znovu využívá, pokud zůstane stejná velikost obrázku.
Pro modely, které přesahují kapacitu jediného čipu, lze metodě from_pretrained předat TensorParallelConfig. Transformer se pak rozděluje mezi čipy už během načítání, takže celý model nemusí být uložen na jednom místě. Podpora ale není u všech pipeline stejná: některé vyžadují doplnění tp_plan.
Ukázky porovnávají kompilovaný režim s režimem eager
Hugging Face předvedl rozdíl mezi oběma režimy na Cloud TPU v6e-4. V ukázkách byly modely rozdělené mezi čtyři čipy a pracovaly s přesností bf16 při velikosti dávky 1. MiniMax-H3 vytvořil video se zvukem o délce 5,2 sekundy a rozlišení 1344 × 768 za 188 sekund při 49 krocích v kompilovaném režimu. V režimu eager stejná ukázka trvala 262 sekund.
U Qwen-Image-2.1 Hugging Face za stejných hardwarových podmínek předvedl tvorbu obrázku 1024 × 1024 při 40 krocích. Kompilovaný režim v této ukázce potřeboval 3,5 sekundy, zatímco eager 8,0 sekund.
Kernels Hub nabízí hotová jádra pro TPU
Další součástí podpory je nabídka hotových TPU jader v Kernels Hubu, odkud se načítají prostřednictvím knihovny Kernels. Jsou napsaná v Pallas; zapojit lze také jádra z knihovny Tokamax. Jedním z uvedených příkladů je ragged-dot pro modely Mixture-of-Experts.
Hugging Face uvádí také výsledky porovnání jádra z Hubu s výchozí implementací na TPU v6e-8. Při velikosti dávky 64 podle něj dosáhlo dekódování u Qwen3-30B-A3B 21,5násobného zrychlení a u Gemma 4 26B-A4B 12,8násobného zrychlení.
Přístup přes Google Cloud a stav podpory
TPU jsou dostupné prostřednictvím Google Cloud, od jednoho čipu až po pody, tedy sestavy s tisíci čipy. Na virtuální stroj vybavený TPU lze backend TorchTPU nainstalovat jedním příkazem.
Podpora zatím neprochází všemi testy. Hugging Face uvádí, že úspěšně prochází 96,5 % testů knihovny Transformers, které se na TPU spouštějí. Tým do repozitáře TorchTPU nahlásil 34 problémů, z nichž 19 už bylo opraveno.
Hugging Face zároveň připravuje pravidelné noční testování Transformers na TPU. Jeho výsledky plánuje zpřístupnit v dashboardu Transformers.
