Google zveřejňuje TorchTPU jako open source a přináší nativní PyTorch na své TPU

Google 20. října 2026 na konferenci PyTorch Conference 2026 oznámil, že TorchTPU, jeho nativní sada nástrojů PyTorch pro tenzorové procesorové jednotky (TPU), je nyní k dispozici jako open source na GitHubu, uvedla společnost na blogu Google Cloud. Software vznikl ve spolupráci s týmem PyTorch ve společnosti Meta a je navržen tak, aby na TPU spouštěl stávající úlohy PyTorch s minimálními úpravami kódu.

Díky tomuto vydání se TPU podle Googlu stávají vedle GPU a CPU nativní možností pro PyTorch. Podporují trénování, nasazování inference i produkční nasazení prostřednictvím nástrojů, jako jsou vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face a LiteRT. Vývojáři mohou modely ladit v eager režimu a poté pomocí torch.compile dosáhnout výkonu na úrovni produkčního nasazení, optimalizovat využití TPU a snížit režii paměti.

Jana van Greunen, ředitelka pro vývoj PyTorch ve společnosti Meta, v oznámení uvedla, že TorchTPU společnosti Meta umožňuje flexibilně vyvíjet hardware společně s jejím hlavním týmem pro systémy strojového učení a zároveň podporovat rozsáhlé úlohy systémů LLM, reklamních a doporučovacích systémů. „Umožňuje nám nasazovat úlohy PyTorch na Google TPU, poskytovat přímou zpětnou vazbu technickým týmům a budovat flexibilní ekosystém, který všichni potřebují,“ řekla.

Integrace Hugging a inferenčního nasazování

Google navázal partnerství s Hugging Face, aby modely s otevřenými vahami na platformě Hugging Face Hub nativně fungovaly v TorchTPU. Ve verzi Transformers v5 se model na TPU nasazuje předáním device_map="tpu", přičemž knihovna Diffusers nativně funguje na TPU a podporuje eager režim, torch.compile i paralelismus tensorů napříč několika akcelerátory. Google uvedl, že také vytvořil vyhrazené CI/CD pipeline pro průběžné noční testování sady Hugging Face Transformers na infrastruktuře TPU. Jeff Boudier, viceprezident pro produkt ve společnosti Hugging Face, uvedl, že tato spolupráce naplňuje závazek platformy nabízet možnost volby hardwaru.

Nativní backend vLLM pro PyTorch v TorchTPU vznikl pro inferenční nasazování ve spolupráci se společností Inferact. Google uvedl, že na hardwaru TPU podporuje průběžné dávkování, kvantizaci a paralelismus napříč několika čipy pomocí vlastních jader Pallas, přičemž zachovává původní modely pracovních procesů, plánovače i API servery, takže stávající inferenční zásobníky nevyžadují žádné změny. Simon Mo, generální ředitel a spoluzakladatel společnosti Inferact, uvedl, že integrace přináší TPU plnohodnotnou podporu vLLM.

Nativní backend PyTorch pro SGLang, vyvinutý ve spolupráci se společností RadixArk, je určen pro strukturované úlohy a úlohy využívající agenty. Podle oznámení SGLang pomocí RadixAttention opakovaně využívá mezipaměť KV napříč požadavky, takže při opakování úvodních částí promptu není nutné výpočty provádět znovu. Verze TorchTPU využívá oddělený backend pro přenos, který zároveň slouží jako vstupně-výstupní backend pro HiCache. Banghua Zhu, technický ředitel a spoluzakladatel společnosti RadixArk, uvedl, že TorchTPU urychlil práci společnosti při postupném zavádění podpory SGLang pro různé typy křemíkových čipů.

Google také ve spolupráci se společnostmi Red Hat a IBM a komunitou čítající více než 500 přispěvatelů a 25 organizací vydává llm-d 1.0 jako projekt v programu Sandbox nadace Cloud Native Computing Foundation. Oznámení popisuje llm-d jako propojení ekosystému pro inferenční nasazování PyTorch s produkčním nasazením založeným na Kubernetes: jeho směrovač s ohledem na prefixovou mezipaměť směruje každý požadavek k replice, která již obsahuje jeho kontext, a odděluje výpočetně náročnou fázi předvyplnění od paměťově náročné fáze dekódování. Verze 1.0 přidává řízení toku pro více uživatelů, vysokou dostupnost směrovače, stabilní rozhraní API v1 a směrování pro agenty zohledňující relaci. V Google Kubernetes Engine může společný koncový bod llm-d běžet na GPU i TPU.

Prof. Sashikumaar Ganesan, zakladatel společnosti Zenteiq, v oznámení uvedl, že GKE Inference Gateway využívající llm-d dosáhl při generování syntetických dat pro tuto společnost více než čtyřnásobné celkové propustnosti inference. Denní kapacita generovaných tokenů se tak zvýšila z 1,47 miliardy na 5,97 miliardy tokenů na každých 1 000 akcelerátorů TPU v6e.

Trénování a orchestrace

TorchTitan, referenční implementace distribuovaného trénování nativní pro PyTorch, nyní obsahuje integraci TorchTPU, kterou Google označuje za připravenou pro produkční nasazení. Je založena na torch.compile, SimpleFSDP v TorchTitanu, specializovaných jádrech Pallas a optimalizacích na úrovni grafu. TorchTitan nabízí recepty TorchTPU pro předtrénování hustých modelů, jako jsou Llama 3 70B a 405B, Qwen 3 8B a 32B a MuseGlimmer 30B, a také modelu Qwen 3 235B typu mixture-of-experts a difuzního transformátoru FLUX.1 12B. Google uvedl, že integraci ověřil až na 8 000 TPU Ironwood (TPU v7x) a že trénování TorchTPU prostřednictvím TorchTitanu dosáhlo u Qwen 3 8B až 52% využití výpočetního výkonu modelu, přičemž škálování bylo téměř lineární až do 8 192 TPU Ironwood.

Google také přispěl backendem pro TPU a kompletním referenčním receptem do VeRL, frameworku pro posilované učení. Recept GRPO, založený na architektuře s odděleným trénovacím a vzorkovacím modulem a předvedený na datových sadách pro matematické uvažování, má vývojářům umožnit spouštět úlohy VeRL na TPU a přitom zachovat stávající kód modelů a algoritmů, který není vázán na konkrétní hardware.

Ray přidává nativní podporu TorchTPU pro orchestraci zpracování dat, trénování, posilovaného učení a nasazování modelů v PyTorch napříč CPU, GPU a TPU v rámci jediného pracovního postupu. Edward Oakes, hlavní architekt ve společnosti Anyscale, uvedl, že firmy spolupracovaly na tom, aby se TPU staly plnohodnotným akcelerátorem v Ray, a dodal: „TorchTPU přímo začleňuje TPU společnosti Google do této oborové standardní platformy.“

Předchozí architektura a dostupnost

TorchTPU byl veřejně představen 7. dubna 2026 v příspěvku na blogu Google Developers, který popisoval návrh „Eager First“ založený na rozhraní PrivateUse1 v PyTorch a režimy eager Debug, Strict a Fused využívající společnou mezipaměť kompilace. Google tehdy uvedl, že režim Fused Eager přináší oproti režimu Strict Eager zvýšení výkonu o 50 % až více než 100 % a že torch.compile v TorchTPU zachycuje grafy FX pomocí Torch Dynamo a převádí je prostřednictvím XLA a StableHLO; vlastní jádra jsou napsána v Pallas a JAX. Podle příspěvku platforma přímo podporuje Distributed Data Parallel, FSDPv2 a DTensor a její architektura počítá s rozdílným průběhem vykonávání napříč procesními skupinami. Mezi položkami plánu vývoje TorchTPU na rok 2026 příspěvek uváděl také spuštění veřejného repozitáře na GitHubu.

Repozitář torch_tpu je nyní veřejný v organizaci google-pytorch na GitHubu a je licencován pod licencí Apache 2.0. Jeho soubor README uvádí vlastní jádra ATen, mezipaměť kompilace pro výpočty XLA, podporu distribuovaného trénování s paralelizací podle dat i tenzorů, integraci s torch.compile, nástroje pro ladění a měření výkonu a příklady pro Llama, Qwen, ResNet, DLRM, minGPT a UNet. Projekt podporuje Python 3.11 a novější verze a balíček torch-tpu 0.2.0 byl zveřejněn na PyPI 19. října 2026; obsahuje balíčky wheel pro Linux x86-64 a CPython 3.11 až 3.14. Uživatelská příručka je dostupná na adrese torchtpu.readthedocs.io a soubor README upozorňuje, že interaktivní notebooky projektu je nutné spouštět z virtuálního počítače TPU; nejjednodušší možností je instance v6e GCE.

Ve stejném oznámení Google uvedl, že se do ekosystému PyTorch připojuje LiteRT, dříve TensorFlow Lite: knihovna LiteRT-Torch převádí model PyTorch jediným litert-torch.convert() voláním a převedené modely lze spouštět na zařízeních se systémy Android, iOS, macOS, Windows a Linux, na webu i v zařízeních IoT. Sada agentů TPU, jejímž cílem je automatizovat zavádění, diagnostiku a ladění výkonu inference u modelů PyTorch, je v rané fázi testování na GitHubu.

Oznámení odkazuje vývojáře na dokumentaci TPU ve službě Google Cloud, kde mohou začít spouštět PyTorch na TPU.