Google publică TorchTPU ca sursă deschisă, aducând PyTorch nativ pe TPU-urile sale

Google a anunțat pe 20 octombrie 2026, la PyTorch Conference 2026, că TorchTPU, stiva sa nativă PyTorch pentru unități de procesare tensorială (TPU-uri), este acum disponibilă ca sursă deschisă pe GitHub, a anunțat compania pe blogul Google Cloud. Software-ul a fost dezvoltat în colaborare cu echipa PyTorch de la Meta și este conceput să ruleze pe TPU-uri sarcinile de lucru PyTorch existente, cu modificări minime ale codului.

Odată cu această lansare, TPU-urile devin, alături de GPU-uri și CPU-uri, o opțiune nativă pentru PyTorch, a declarat Google, acoperind antrenarea, servirea și implementarea prin instrumente precum vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face și LiteRT. Dezvoltatorii pot depana modelele în modul eager, apoi pot obține performanțe de nivelul celor necesare în producție cu ajutorul torch.compile, pentru a optimiza utilizarea TPU-urilor și a reduce memoria suplimentară necesară.

Jana van Greunen, directoarea departamentului de inginerie PyTorch de la Meta, a declarat în anunț că TorchTPU îi permite companiei Meta să dezvolte în comun componente hardware, cu flexibilitate, alături de echipa sa centrală de sisteme de învățare automată, susținând totodată sarcinile de lucru pentru sistemele majore LLM, de publicitate și de recomandare. „Ne permite să implementăm sarcinile noastre de lucru PyTorch pe TPU-urile Google, să oferim feedback direct echipelor de inginerie și să construim flexibilitatea ecosistemului de care are nevoie toată lumea”, a spus ea.

Integrări pentru Hugging și servire

Google a colaborat cu Hugging Face pentru ca modelele cu ponderi deschise de pe Hugging Face Hub să ruleze nativ pe TorchTPU. În Transformers v5, implementarea unui model pe TPU-uri se face prin transmiterea device_map="tpu", iar biblioteca Diffusers rulează nativ pe TPU-uri și acceptă modul eager, torch.compile și paralelismul tensorial pe mai mulți acceleratori. Google a declarat că a creat și fluxuri dedicate CI/CD pentru testarea nocturnă continuă a suitei Hugging Face Transformers pe infrastructura TPU. Jeff Boudier, vicepreședinte de produs la Hugging Face, a spus că această colaborare pune în practică angajamentul platformei față de libertatea de a alege hardware-ul.

Pentru servire, backendul vLLM nativ PyTorch pentru TorchTPU a fost dezvoltat în colaborare cu Inferact. Google a declarat că acesta acceptă gruparea continuă a loturilor, cuantizarea și paralelismul pe mai multe cipuri pe hardware TPU, folosind nuclee Pallas personalizate, păstrând în același timp modelele de worker, planificatoarele și serverele API din proiectul-sursă, astfel încât stivele de inferență existente să nu necesite modificări. Simon Mo, CEO și cofondator al Inferact, a spus că integrarea oferă TPU-urilor suport de prim rang pentru vLLM.

Un backend PyTorch nativ pentru SGLang, dezvoltat împreună cu RadixArk, este destinat sarcinilor de lucru structurate și agentice. Potrivit anunțului, SGLang folosește RadixAttention pentru a reutiliza memoria cache KV între solicitări, astfel încât prefixele de prompt repetate să nu mai necesite recalculare, iar versiunea TorchTPU utilizează un backend de transfer cu dezagregare, care servește și drept backend I/O pentru HiCache. Banghua Zhu, director tehnic și cofondator al RadixArk, a declarat că TorchTPU a accelerat activitatea companiei în timp ce aceasta extinde suportul SGLang la diferite tipuri de siliciu.

Google lansează, de asemenea, llm-d 1.0 împreună cu Red Hat, IBM și o comunitate de peste 500 de colaboratori și 25 de organizații, ca proiect Sandbox al Cloud Native Computing Foundation. Anunțul descrie llm-d drept o punte între ecosistemul de servire PyTorch și implementarea în producție bazată pe Kubernetes: routerul său folosește rutarea care ține cont de memoria cache a prefixelor pentru a trimite fiecare solicitare către replica ce deține deja contextul acesteia și separă etapa de preumplere, solicitantă din punct de vedere al calculului, de etapa de decodare, solicitantă din punct de vedere al memoriei. Versiunea 1.0 adaugă controlul fluxului pentru mai mulți locatari, disponibilitate ridicată a routerului, API-uri v1 stabile și rutare pentru agenți care ține cont de sesiune; în plus, pe Google Kubernetes Engine, un punct final llm-d comun poate rula fie pe GPU-uri, fie pe TPU-uri.

Prof. Sashikumaar Ganesan, fondatorul Zenteiq, a declarat în anunț că GKE Inference Gateway, bazat pe llm-d, a oferit un debit agregat de inferență de peste 4 ori mai mare pentru generarea de date sintetice a companiei, mărind capacitatea zilnică de generare de la 1,47 miliarde la 5,97 miliarde de tokenuri pentru fiecare 1.000 de acceleratoare TPU v6e.

Antrenare și orchestrare

TorchTitan, implementarea de referință nativă PyTorch pentru antrenarea distribuită, include acum ceea ce Google descrie drept o integrare TorchTPU de nivelul cerut în producție, construită pe torch.compile, SimpleFSDP din TorchTitan, nuclee Pallas specializate și optimizări la nivel de graf. TorchTitan oferă rețete TorchTPU pentru preantrenarea unor modele dense precum Llama 3 70B și 405B, Qwen 3 8B și 32B și MuseGlimmer 30B, precum și a modelului Mixture of Experts Qwen 3 235B și a transformatorului de difuzie FLUX.1 12B. Google a declarat că integrarea a fost validată pe până la 8.000 de TPU-uri Ironwood (TPU v7x) și că antrenarea TorchTPU prin TorchTitan a atins o utilizare a FLOP-urilor modelului de până la 52% pentru Qwen 3 8B, cu o scalare aproape liniară pe până la 8.192 de TPU-uri Ironwood.

Google a integrat, de asemenea, în VeRL, un cadru de învățare prin consolidare, un backend TPU și o rețetă de referință completă, de la un capăt la altul. Rețeta GRPO, construită pe o arhitectură cu antrenor și eșantionator dezagregate și demonstrată pe seturi de date pentru raționament matematic, este concepută să le permită dezvoltatorilor să ruleze pe TPU-uri sarcini de lucru VeRL, păstrându-și codul existent pentru modele și algoritmi, independent de hardware.

Ray adaugă suport nativ pentru TorchTPU, permițând orchestrarea prelucrării datelor, a antrenării, a învățării prin consolidare și a servirii modelelor PyTorch pe procesoare, GPU-uri și TPU-uri, într-un singur flux de lucru. Edward Oakes, arhitect-șef la Anyscale, a declarat că cele două companii au colaborat pentru ca TPU-urile să devină acceleratoare de prim rang în Ray, adăugând: „TorchTPU integrează direct TPU-urile Google în această stivă standard în industrie.”

Arhitectura anterioară și disponibilitatea

TorchTPU a fost prezentat public pe 7 aprilie 2026, într-o postare pe blogul Google Developers care descria un design „Eager First”, bazat pe interfața PrivateUse1 din PyTorch, cu modurile eager Debug, Strict și Fused, susținute de un cache de compilare partajat. La momentul respectiv, Google a afirmat că Fused Eager oferă o creștere a performanței de la 50% până la peste 100% față de Strict Eager și că torch.compile pe TorchTPU captează grafuri FX cu Torch Dynamo și le transformă prin XLA și StableHLO, folosind nuclee personalizate scrise în Pallas și JAX. Potrivit postării, stiva oferă din start suport pentru Distributed Data Parallel, FSDPv2 și DTensor și este concepută pentru a gestiona execuția divergentă între ranguri. Postarea a inclus, de asemenea, lansarea unui depozit public pe GitHub printre obiectivele foii de parcurs TorchTPU pentru 2026.

Depozitul torch_tpu este acum public în cadrul organizației GitHub google-pytorch și este licențiat sub Apache 2.0. Fișierul README enumeră nuclee ATen personalizate, un cache de compilare pentru calculele XLA, suport pentru antrenare distribuită cu paralelism pe date și pe tensori, integrare cu torch.compile, utilitare de depanare și evaluare a performanței, precum și exemple pentru Llama, Qwen, ResNet, DLRM, minGPT și UNet. Proiectul este compatibil cu Python 3.11 și versiunile ulterioare, iar pachetul torch-tpu 0.2.0 a fost publicat pe PyPI la 19 octombrie 2026, cu pachete wheel pentru Linux x86-64 și CPython de la versiunea 3.11 până la 3.14. Ghidul utilizatorului este publicat la torchtpu.readthedocs.io, iar fișierul README precizează că notebookurile interactive ale proiectului trebuie rulate de pe o mașină virtuală TPU, cea mai simplă opțiune fiind o instanță GCE v6e.

În același anunț, Google a declarat că LiteRT, cunoscut anterior ca TensorFlow Lite, se alătură ecosistemului PyTorch: biblioteca LiteRT-Torch convertește un model PyTorch printr-un singur litert-torch.convert() apel, iar modelele convertite rulează pe dispozitive Android, iOS, macOS, Windows, Linux, web și IoT. O suită de agenți TPU, concepută pentru a automatiza integrarea, diagnosticarea și optimizarea performanței inferenței pentru modelele PyTorch, se află în faza de previzualizare timpurie pe GitHub.

Anunțul îi îndrumă pe dezvoltatori către documentația TPU din Google Cloud pentru a începe să ruleze PyTorch pe TPU-uri.