Google ogłosiło TorchTPU 20 października 2026 r. To otwartoźródłowy backend, który pozwala uruchamiać PyTorch na akceleratorach TPU Google. W podstawowym przykładzie wystarczy wskazać urządzenie tpu zamiast cuda; zadania rozproszone wymagają dodatkowej konfiguracji właściwej dla TPU.
Dla programisty oznacza to prostą ścieżkę do sprawdzenia, czy kod ruszy na TPU, ale nie obietnicę, że każda aplikacja zadziała bez zmian albo od razu osiągnie maksymalną wydajność. TorchTPU rozwijano we współpracy z zespołem PyTorch w Meta.
Podstawowe uruchomienie PyTorch na TPU
Dokumentacja Google podaje instalację poleceniem pip install torch_tpu i wybór urządzenia przez torch.device("tpu"). W kodzie korzystającym z .to(...) przykładowa zmiana wygląda tak: .to("cuda") na .to("tpu"). Dokumentacja Google Cloud dotycząca PyTorch na TPU opisuje także wykonywanie w trybie eager, domyślne dla TorchTPU, oraz kompilowanie modelu przez torch.compile.
| Rodzaj zadania | Wybór urządzenia lub backendu | Dodatkowa konfiguracja uruchomieniowa |
| Podstawowe | torch.device("tpu") lub .to("tpu") zamiast urządzenia CUDA | Instalacja torch_tpu; wykonanie w trybie eager albo kompilacja przez torch.compile |
| Rozproszone | Backend grupy procesów tpu_dist; jeden proces na urządzenie | Konfiguracja topologii TPU i slice-buildera lub użycie udokumentowanego narzędzia pomocniczego |
Różnica staje się istotna, gdy model działa na wielu urządzeniach: sama zmiana wskazania urządzenia nie konfiguruje uruchomienia rozproszonego. W takim przypadku TorchTPU wykorzystuje backend tpu_dist, a konfiguracja startowa musi uwzględniać topologię i slice-builder TPU.
Zgodność zależy od modelu i sposobu użycia
Zespół projektu podaje obsługę około 450 operatorów ATen, z wyłączeniem delegatów strukturalnych lub złożonych. W jego testach 4 821 z 4 917 wybranych modeli Tier 1 uruchomiło się bez awarii, a 638 z 657 reprezentatywnych modeli Tier 2 osiągnęło zgodność numeryczną z wynikiem referencyjnym na GPU. W osobnym teście poprawności dla modeli językowych z ograniczoną dynamiką kształtów wynik wyniósł 834 na 839.
To wyniki projektu TorchTPU dla określonych zbiorów i testów, a nie gwarancja zgodności każdego modelu ani jednolity benchmark GPU kontra TPU. Wpis techniczny zespołu opisuje też przykład przeniesienia Muse Glimmer 30B, w którym w testowanej implementacji zastąpiono FlexAttention operacją F.scaled_dot_product_attention, ponieważ FlexAttention nie było wtedy obsługiwane w TorchTPU. Pokazuje to, że niektóre modele mogą wymagać zmian w konkretnych operacjach.
Projekt PyTorch opisał również trening modelu Qwen3 MoE 235B w konfiguracjach do 8 192 układów TPU v7. W największej konfiguracji, przy długości sekwencji 4 096 i batchu 262 144, zespół podał 2 104 tokeny na sekundę na TensorCore, 27,1% MFU i 31,2 sekundy na krok. Są to wyniki dla tej konfiguracji treningowej; nie stanowią porównania wydajności z GPU. Techniczny wpis projektu PyTorch o TorchTPU opisuje zakres tych testów i ich ustawienia.
Wczytywanie wag zapisanych po treningu na GPU
Zgodne wagi modelu można wczytać do modelu PyTorch przygotowanego do pracy na TPU. state_dict przechowuje parametry modelu, a torch.load umożliwia mapowanie danych checkpointu na wybrane urządzenie. Dokumentacja PyTorch opisuje zapisywanie i wczytywanie modeli. Przy wznawianiu treningu wagi modelu należy odróżnić od stanu optymalizatora i innych elementów treningu, które są odrębnymi składnikami checkpointu.
