Google otwiera TorchTPU. PyTorch działa natywnie na chipach TPU

Czym jest TorchTPU

TorchTPU to warstwa, która pozwala PyTorchowi traktować TPU jak zwykłe urządzenie obliczeniowe, obok procesorów i kart graficznych. Repozytorium zawiera implementację operacji PyTorch (tzw. kerneli ATen) dla TPU, pamięć podręczną kompilacji, obsługę równoległości danych i tensorów oraz przykłady dla modeli Llama, Qwen, ResNet, DLRM i UNet.

Do tej pory główną drogą uruchamiania PyTorcha na TPU był projekt PyTorch/XLA, który wymagał od programistów dostosowania kodu do modelu działania kompilatora XLA. Google przedstawia TorchTPU jako rozwiązanie natywne, z debugowaniem w trybie eager i przejściem do torch.compile, gdy model jest gotowy do produkcji.

Integracje z vLLM, SGLang i Hugging Face

Dla zespołów, które uruchamiają modele językowe, ważniejsze od samego backendu są gotowe integracje. Google przygotował natywne backendy PyTorch dla dwóch popularnych silników inferencji. Dla vLLM powstał we współpracy z firmą Inferact i obsługuje m.in. ciągłe grupowanie zapytań, kwantyzację i pracę na wielu chipach. Dla SGLang przygotowano go razem z RadixArk.

W Hugging Face Transformers v5 wiele modeli z Hub działa na TPU po wskazaniu urządzenia parametrem device_map="tpu". Na TPU działa też biblioteka Diffusers do modeli generujących obrazy. Według Google zestaw testów Transformers jest codziennie uruchamiany na infrastrukturze TPU. Lista integracji obejmuje również llm-d 1.0, TorchTitan, VeRL, Ray i TorchRec.

Wydajność według Google

Google podał dwa przykłady wyników. Oba pochodzą od producenta i nie zostały niezależnie zweryfikowane.

• Uczenie modelu Qwen 3 8B w TorchTitan - do 52 proc. wykorzystania mocy obliczeniowej (Model FLOPs Utilization), z niemal liniowym skalowaniem do 8192 chipów Ironwood (TPU v7x).
• Inferencja w llm-d u firmy Zenteiq na TPU v6e - ponad czterokrotny wzrost przepustowości, z 1,47 mld do 5,97 mld wygenerowanych tokenów dziennie na 1000 akceleratorów.
• Meta, współtwórca TorchTPU, korzysta z niego przy modelach językowych, reklamach i systemach rekomendacji.

TPU agent suite i LiteRT

Równolegle Google udostępnił w wersji preview TPU agent suite. To zestaw agentów AI, które przez wtyczki współpracują z narzędziami do programowania i pomagają przenieść model na TPU, diagnozować błędy oraz stroić wydajność inferencji. Google nie podał listy obsługiwanych asystentów kodowania.

Drugą zmianą jest dołączenie LiteRT, dawnego TensorFlow Lite, do ekosystemu PyTorch. Narzędzie LiteRT-Torch konwertuje model PyTorch jednym wywołaniem funkcji, a modele z Transformers można eksportować do uruchamiania na telefonach, komputerach i urządzeniach IoT.

Google i Amazon przeciw przewadze NVIDII

Ogłoszenie zbiega się z premierą TorchNeuron, czyli natywnego backendu PyTorch dla chipów AWS Trainium, który Amazon pokazał 19 października. Obie firmy celują w ten sam problem. Większość kodu AI powstaje w PyTorchu z myślą o kartach NVIDII i środowisku CUDA, więc przeniesienie go na inny sprzęt wymagało dotąd pracy, której zespoły chciały uniknąć.

W naszej ocenie dla zespołów ML oznacza to realną możliwość negocjowania kosztów mocy obliczeniowej. Jeśli ten sam kod da się uruchomić na kartach NVIDII, TPU i Trainium przy niewielkich zmianach, wybór chmury mniej zależy od dostępności konkretnego sprzętu. Deklaracje producentów trzeba jednak sprawdzić na własnych modelach.

Co to oznacza dla zespołów w Polsce

TPU są dostępne wyłącznie w Google Cloud, więc TorchTPU dotyczy firm i uczelni, które korzystają z tej chmury lub ją rozważają. Google Cloud ma region europe-central2 w Warszawie, ale dostępność konkretnych generacji TPU różni się między regionami. Przed planowaniem projektu trzeba sprawdzić w dokumentacji Google Cloud, gdzie dostępny jest potrzebny typ chipu.

Najprostszym testem jest uruchomienie istniejącego modelu z Transformers na maszynie wirtualnej z TPU v6e, którą zaleca dokumentacja projektu. Porównanie kosztu godziny pracy i przepustowości z obecną konfiguracją na kartach graficznych pokaże, czy migracja ma sens finansowy.

Ograniczenia

Projekt ma numer wersji 0.2.0, co wskazuje na wczesny etap rozwoju. Google nie podał w ogłoszeniu pełnej listy obsługiwanych generacji TPU ani wersji PyTorcha, a TPU agent suite jest dopiero w wersji preview. Wyniki wydajności pochodzą od Google i jego klienta. Nie wiadomo też, jak szybko TorchTPU obejmie mniej popularne operacje i własne kernele, które zespoły napisały wcześniej pod CUDA.