Google udostępnia TorchTPU jako oprogramowanie open source, wprowadzając natywną obsługę PyTorch na swoich TPU

20 października 2026 r. podczas konferencji PyTorch Conference 2026 Google ogłosiło, że TorchTPU — natywny stos PyTorch dla jednostek przetwarzania tensorowego (TPU) — jest już dostępny jako oprogramowanie open source w serwisie GitHub, poinformowała firma na blogu Google Cloud. Oprogramowanie opracowano wspólnie z zespołem PyTorch w Meta; ma ono umożliwiać uruchamianie istniejących obciążeń PyTorch na TPU przy minimalnych zmianach w kodzie.

Dzięki tej wersji TPU dołączają do GPU i CPU jako natywna opcja dla PyTorch — poinformowało Google. Obsługiwane są zadania związane z trenowaniem, obsługą inferencji i wdrażaniem, realizowane za pomocą narzędzi takich jak vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face i LiteRT. Programiści mogą debugować modele w trybie eager, a następnie osiągać wydajność na poziomie produkcyjnym dzięki torch.compile, aby zoptymalizować wykorzystanie TPU i zmniejszyć narzut pamięciowy.

Jana van Greunen, dyrektorka ds. inżynierii PyTorch w Meta, powiedziała w komunikacie, że TorchTPU pozwala Meta elastycznie współtworzyć sprzęt z zespołem odpowiedzialnym za podstawowe systemy uczenia maszynowego, a zarazem obsługiwać główne obciążenia związane z LLM, systemami reklamowymi i rekomendacyjnymi. „Dzięki temu możemy wdrażać nasze obciążenia PyTorch na TPU Google, przekazywać bezpośrednie informacje zwrotne zespołom inżynierskim i rozwijać elastyczny ekosystem, którego wszyscy potrzebują” — powiedziała.

Integracje z Hugging i systemami obsługującymi inferencję

Google nawiązało współpracę z Hugging Face, aby modele z otwartymi wagami dostępne w Hugging Face Hub działały natywnie w TorchTPU. W Transformers v5 wdrożenie modelu na TPU odbywa się przez przekazanie device_map="tpu", a biblioteka Diffusers działa natywnie na TPU i obsługuje tryb eager, torch.compile oraz równoległość tensorową obejmującą wiele akceleratorów. Google poinformowało również o przygotowaniu specjalnych potoków CI/CD do ciągłego, nocnego testowania pakietu Hugging Face Transformers na infrastrukturze TPU. Jeff Boudier, wiceprezes ds. produktu w Hugging Face, powiedział, że ta współpraca przekłada na praktykę zobowiązanie platformy do zapewniania użytkownikom wyboru sprzętu.

Backend vLLM dla natywnego PyTorch w TorchTPU, przeznaczony do obsługi inferencji, opracowano we współpracy z Inferact. Google poinformowało, że obsługuje on ciągłe grupowanie zadań, kwantyzację i równoległość wieloukładową na sprzęcie TPU z użyciem niestandardowych jąder Pallas, a jednocześnie zachowuje oryginalne modele procesów roboczych, harmonogramy i serwery API, dzięki czemu istniejące stosy inferencyjne nie wymagają zmian. Simon Mo, dyrektor generalny i współzałożyciel Inferact, powiedział, że integracja zapewnia TPU pełnoprawną obsługę vLLM.

Natywny backend PyTorch dla SGLang, opracowany we współpracy z RadixArk, jest przeznaczony do zadań o ustrukturyzowanym charakterze i obciążeń agentowych. Jak podano w komunikacie, SGLang wykorzystuje RadixAttention do ponownego używania pamięci podręcznej KV w kolejnych żądaniach, dzięki czemu powtarzające się prefiksy promptów nie wymagają ponownych obliczeń. Wersja TorchTPU korzysta z backendu transferu z dezagregacją, który pełni również funkcję backendu wejścia/wyjścia dla HiCache. Banghua Zhu, dyrektor ds. technologii i współzałożyciel RadixArk, powiedział, że TorchTPU przyspieszył prace firmy nad wdrażaniem obsługi SGLang na różnych typach układów krzemowych.

Google udostępnia również llm-d 1.0 we współpracy z Red Hat, IBM oraz społecznością liczącą ponad 500 współtwórców i 25 organizacji, jako projekt w programie Sandbox Cloud Native Computing Foundation. W komunikacie llm-d opisano jako rozwiązanie łączące ekosystem obsługi inferencji PyTorch z wdrożeniami produkcyjnymi opartymi na Kubernetes: jego router kieruje każde żądanie do repliki, która ma już zapisany jego kontekst, wykorzystując routing uwzględniający pamięć podręczną prefiksów, a także oddziela wymagające dużej mocy obliczeniowej wstępne przetwarzanie od wymagającego dużej ilości pamięci dekodowania. Wersja 1.0 wprowadza kontrolę przepływu dla wielu dzierżawców, wysoką dostępność routera, stabilne interfejsy API v1 oraz routing uwzględniający sesje agentów. Na Google Kubernetes Engine wspólny punkt końcowy llm-d może działać zarówno na GPU, jak i na TPU.

Prof. Sashikumaar Ganesan, założyciel Zenteiq, powiedział w komunikacie, że GKE Inference Gateway oparty na llm-d zapewnił ponad 4× większą łączną przepustowość inferencji na potrzeby generowania danych syntetycznych przez firmę, zwiększając dzienną liczbę generowanych tokenów z 1,47 mld do 5,97 mld na 1000 akceleratorów TPU v6e.

Trenowanie i orkiestracja

TorchTitan, natywna dla PyTorch implementacja referencyjna do trenowania rozproszonego, obejmuje teraz — jak opisuje to Google — integrację TorchTPU klasy produkcyjnej, opartą na torch.compile, SimpleFSDP w TorchTitan, wyspecjalizowanych jądrach Pallas i optymalizacjach na poziomie grafu. TorchTitan udostępnia przepisy TorchTPU do wstępnego trenowania gęstych modeli, takich jak Llama 3 70B i 405B, Qwen 3 8B i 32B oraz MuseGlimmer 30B, a także modelu Qwen 3 235B typu mieszanka ekspertów i transformera dyfuzyjnego FLUX.1 12B. Google poinformowało, że integrację zweryfikowano na maksymalnie 8000 TPU Ironwood (TPU v7x), a trenowanie TorchTPU z użyciem TorchTitan osiągnęło do 52% wykorzystania mocy obliczeniowej modelu w przypadku Qwen 3 8B, przy niemal liniowym skalowaniu na maksymalnie 8192 TPU Ironwood.

Google przekazało również do VeRL — frameworka uczenia przez wzmacnianie — backend TPU oraz referencyjny przepis obejmujący cały proces. Przepis GRPO, oparty na architekturze z oddzielonym trenerem i modułem próbkowania oraz zaprezentowany na zbiorach danych do rozumowania matematycznego, ma umożliwić programistom uruchamianie obciążeń VeRL na TPU przy zachowaniu istniejącego kodu modeli i algorytmów, niezależnego od sprzętu.

Ray zapewnia natywną obsługę TorchTPU, umożliwiającą koordynowanie przetwarzania danych, trenowania, uczenia ze wzmocnieniem i udostępniania modeli PyTorch na procesorach CPU, GPU i TPU w ramach jednego przepływu pracy. Edward Oakes, główny architekt w Anyscale, powiedział, że firmy pracowały nad tym, by TPU stały się akceleratorami pierwszej klasy w Ray, dodając: „TorchTPU integruje TPU Google bezpośrednio z tym standardowym w branży stosem”.

Wcześniejsza architektura i dostępność

TorchTPU zaprezentowano publicznie 7 kwietnia 2026 r. we wpisie na blogu Google Developers, w którym opisano architekturę „Eager First” opartą na interfejsie PrivateUse1 PyTorch, z trybami eager Debug, Strict i Fused korzystającymi ze wspólnej pamięci podręcznej kompilacji. Google poinformowało wówczas, że Fused Eager zapewnia wzrost wydajności od 50% do ponad 100% w porównaniu ze Strict Eager, a torch.compile w TorchTPU przechwytuje grafy FX za pomocą Torch Dynamo i przekształca je z użyciem XLA i StableHLO; niestandardowe jądra są pisane w Pallas i JAX. We wpisie podano, że stos obsługuje od razu Distributed Data Parallel, FSDPv2 i DTensor oraz został zaprojektowany tak, by radzić sobie z rozbieżnym wykonywaniem zadań na poszczególnych rangach. Wśród elementów planu rozwoju TorchTPU na 2026 r. wymieniono również uruchomienie publicznego repozytorium w serwisie GitHub.

Repozytorium torch_tpu jest już publicznie dostępne w organizacji google-pytorch w serwisie GitHub i udostępniane na licencji Apache 2.0. W pliku README wymieniono niestandardowe jądra ATen, pamięć podręczną kompilacji obliczeń XLA, obsługę trenowania rozproszonego z równoległością danych i tensorów, integrację z torch.compile, narzędzia do debugowania i testów wydajnościowych oraz przykłady dla modeli Llama, Qwen, ResNet, DLRM, minGPT i UNet. Projekt obsługuje Python 3.11 i nowsze wersje, a pakiet torch-tpu 0.2.0 opublikowano w serwisie PyPI 19 października 2026 r.; zawiera on pakiety binarne dla systemu Linux x86-64, przeznaczone dla CPython od wersji 3.11 do 3.14. Podręcznik użytkownika jest dostępny pod adresem torchtpu.readthedocs.io, a w pliku README zaznaczono, że interaktywne notatniki projektu muszą być uruchamiane na maszynie wirtualnej TPU; najłatwiej skorzystać w tym celu z instancji GCE v6e.

W tym samym komunikacie Google poinformowało, że LiteRT, wcześniej znane jako TensorFlow Lite, dołącza do ekosystemu PyTorch: biblioteka LiteRT-Torch konwertuje model PyTorch za pomocą jednego litert-torch.convert() wywołania, a przekonwertowane modele działają na urządzeniach z systemami Android, iOS, macOS, Windows i Linux, w internecie oraz w urządzeniach IoT. Zestaw agentów TPU, który ma automatyzować wdrażanie, diagnostykę i dostrajanie wydajności wnioskowania modeli PyTorch, jest dostępny w serwisie GitHub we wczesnej wersji zapoznawczej.

W komunikacie skierowano deweloperów do dokumentacji TPU w Google Cloud, aby mogli rozpocząć uruchamianie PyTorch na TPU.