Hugging Face ogłasza natywne wsparcie dla TPU, układów akceleracyjnych zaprojektowanych przez Google, w swoim ekosystemie. Dzięki backendowi TorchTPU dla PyTorcha, Transformers, Diffusers i Kernels działają bezpośrednio na tych procesorach, jak wynika z wpisu opublikowanego przez firmę.
TPU, czyli Tensor Processing Unit, to układ zaprojektowany do przyspieszania obliczeń macierzowych, które zajmują większość czasu działania sieci neuronowych. Google rozwija te układy od około dziesięciu lat i wykorzystuje je do trenowania i uruchamiania własnych modeli, w tym Gemini. Są one dostępne w Google Cloud, od pojedynczych układów po zestawy liczące kilka tysięcy.
Do tej pory ekosystem oprogramowania TPU był budowany wokół TensorFlow i JAX, a użytkownicy PyTorcha musieli wkładać więcej wysiłku. Aby uruchomić PyTorch na TPU, trzeba było korzystać z PyTorch/XLA (torch_xla), rozwiązania, które wpis opisuje jako działające, ale niewygodne: tensory były ewaluowane leniwie, przez co kod nie zawsze wykonywał się w oczekiwanym momencie, a drobne zmiany mogły wywołać ponowną kompilację grafu. Hugging Face opracował Optimum TPU, aby umożliwić inferencję i trening swoich modeli na TPU. Dzięki TorchTPU to narzędzie nie jest już potrzebne do uruchamiania modeli z Transformers i Diffusers.
TorchTPU, natywny backend dla PyTorcha
TorchTPU dodaje TPU jako nowy backend PyTorcha, bez modyfikacji zwykłego kodu PyTorcha. Backend obejmuje zoptymalizowane dla TPU jądra Pallas, wsparcie dla treningu rozproszonego oraz integrację z torch.compile dla użytkowników szukających większej prędkości. Jego instalacja na maszynie wirtualnej TPU odbywa się za pomocą jednego polecenia: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. Wpis precyzuje, że ten indeks dostarcza wersję CPU torch, lżejszą niż wersja CUDA, którą pip pobiera domyślnie i której TPU nie potrzebuje.
Transformers, Diffusers i Kernels na TPU
W Transformers v5, którego celem jest bycie referencyjnym dla definicji modeli, aby działały wszędzie, załadowanie modelu na TPU odbywa się poprzez parametr device_map= »tpu ». Wpis ilustruje ten krok na przykładzie modelu Qwen/Qwen3.5-9B. Domyślnie model działa w trybie eager, ale może też korzystać z torch.compile lub konfiguracji rozproszonej.
W Diffusers TPU jest teraz urządzeniem pierwszego wyboru. Tryb eager, odpowiedni do prototypowania, wymaga jedynie wywołania pipe.to(« tpu »), podczas gdy tryb skompilowany, przeznaczony do wdrożenia, dodaje wywołanie kompilacji. Dla modeli zbyt dużych dla pojedynczego układu, TensorParallelConfig przekazany do from_pretred rozdziela transformer między kilka układów przy ładowaniu. Wpis prezentuje przykłady wykonane na Cloud TPU v6e-4, z modelami MiniMax-H3 dla wideo i audio oraz Qwen-Image-2.1.
Biblioteka Kernels zapewnia dostęp do gotowych do użycia jąder TPU, napisanych w Pallas, które oferuje niskopoziomową kontrolę nad sprzętem. Jądra z biblioteki Tokamax również mogą być zintegrowane. Według wpisu, dla batcha wielkości 64, dekodowanie jest 21,5 razy szybsze na Qwen3-30B-A3B i 12,8 razy szybsze na Gemma 4 26B-A4B w porównaniu z implementacją domyślną, przy użyciu TPU v6e-8.
Ciągłe testy
Hugging Face informuje, że stale wykonuje zestaw testów Transformers na TPU i pracuje nad dodaniem nocnej integracji ciągłej na tych układach, której wyniki byłyby publikowane na dedykowanym panelu. Według wpisu, 96,5% testów uruchomionych na TPU kończy się sukcesem. Każda porażka skutkuje skryptem reprodukcyjnym i zgłoszeniem w repozytorium TorchTPU: zgłoszono 34 problemy, z których 19 jest już naprawionych. Testy te ujawniły również błędy w Transformers, m.in. w StaticCache i w planach równoległości tensorowej Qwen3, które zostały naprawione u źródła. W przypadku problemów użytkownicy są zachęcani do otwierania zgłoszeń.
