Google rende TorchTPU open source, portando PyTorch nativo sulle sue TPU

Il 20 ottobre 2026, in occasione della PyTorch Conference 2026, Google ha annunciato che TorchTPU, il suo stack nativo di PyTorch per le Tensor Processing Unit (TPU), è ora disponibile come progetto open source su GitHub, ha dichiarato l’azienda sul blog di Google Cloud. Il software è stato sviluppato insieme al team PyTorch di Meta ed è progettato per eseguire sulle TPU i carichi di lavoro PyTorch esistenti con modifiche minime al codice.

Con questa uscita, le TPU diventano un’opzione nativa per PyTorch, al pari di GPU e CPU, ha dichiarato Google, per addestramento, inferenza e distribuzione tramite strumenti tra cui vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face e LiteRT. Gli sviluppatori possono eseguire il debug dei modelli in modalità eager e poi raggiungere prestazioni di livello produttivo con torch.compile, ottimizzando l’utilizzo delle TPU e riducendo il sovraccarico di memoria.

Jana van Greunen, direttrice dell’ingegneria PyTorch presso Meta, ha dichiarato nell’annuncio che TorchTPU consente a Meta di sviluppare in modo flessibile l’hardware insieme al proprio team centrale di sistemi di apprendimento automatico, supportando al contempo i principali carichi di lavoro dei sistemi LLM, pubblicitari e di raccomandazione. «Ci permette di distribuire i nostri carichi di lavoro PyTorch sulle TPU di Google, fornire un riscontro diretto ai team di ingegneria e realizzare la flessibilità dell’ecosistema di cui tutti hanno bisogno», ha affermato.

Integrazioni con Hugging e per l’inferenza

Google ha collaborato con Hugging Face affinché i modelli open-weight disponibili sull’Hub di Hugging Face funzionino in modo nativo su TorchTPU. In Transformers v5, per distribuire un modello sulle TPU basta passare device_map="tpu", e la libreria Diffusers funziona in modo nativo sulle TPU, con supporto per la modalità eager, torch.compile e il parallelismo tensoriale su più acceleratori. Google ha dichiarato di aver inoltre creato pipeline CI/CD dedicate per eseguire ogni notte test continui della suite Hugging Face Transformers sull’infrastruttura TPU. Jeff Boudier, vicepresidente del prodotto presso Hugging Face, ha affermato che la collaborazione traduce in pratica l’impegno della piattaforma a offrire libertà di scelta dell’hardware.

Per l’inferenza, il backend nativo PyTorch di vLLM per TorchTPU è stato sviluppato in collaborazione con Inferact. Google ha dichiarato che supporta il batching continuo, la quantizzazione e il parallelismo multi-chip sull’hardware TPU tramite kernel Pallas personalizzati, mantenendo al contempo i modelli di worker, gli scheduler e i server API del progetto upstream, così che gli stack di inferenza esistenti non richiedano modifiche. Simon Mo, CEO e cofondatore di Inferact, ha affermato che l’integrazione offre alle TPU un supporto di primo livello per vLLM.

Un backend nativo PyTorch per SGLang, sviluppato con RadixArk, è destinato ai carichi di lavoro strutturati e basati su agenti. Secondo l’annuncio, SGLang usa RadixAttention per riutilizzare la cache KV tra le richieste, così che i prefissi dei prompt ripetuti non debbano essere ricalcolati; la versione per TorchTPU impiega inoltre un backend di trasferimento disaggregato, che funge anche da backend I/O per HiCache. Banghua Zhu, CTO e cofondatore di RadixArk, ha dichiarato che TorchTPU ha accelerato il lavoro dell’azienda mentre estende il supporto a SGLang a diversi tipi di silicio.

Google rilascia inoltre llm-d 1.0 insieme a Red Hat, IBM e a una comunità di oltre 500 collaboratori e 25 organizzazioni, come progetto Sandbox della Cloud Native Computing Foundation. L’annuncio descrive llm-d come un collegamento tra l’ecosistema PyTorch per l’inferenza e la distribuzione in produzione basata su Kubernetes: il suo router usa un instradamento che tiene conto della cache dei prefissi per inviare ogni richiesta alla replica che già contiene il contesto pertinente e separa la fase di prefill, che richiede molta potenza di calcolo, dalla fase di decode, che richiede molta memoria. La versione 1.0 aggiunge il controllo del flusso multi-tenant, l’alta disponibilità del router, API v1 stabili e l’instradamento per agenti basato sulle sessioni; inoltre, su Google Kubernetes Engine, un endpoint llm-d comune può funzionare sia su GPU sia su TPU.

Il professor Sashikumaar Ganesan, fondatore di Zenteiq, ha dichiarato nell’annuncio che GKE Inference Gateway, basato su llm-d, ha fornito un throughput aggregato dell’inferenza superiore di oltre 4 volte per la generazione di dati sintetici dell’azienda, aumentando la capacità giornaliera di generazione da 1,47 miliardi a 5,97 miliardi di token per 1.000 acceleratori TPU v6e.

Addestramento e orchestrazione

TorchTitan, l’implementazione di riferimento nativa per PyTorch per l’addestramento distribuito, ora include quella che Google descrive come un’integrazione TorchTPU pronta per la produzione, basata su torch.compile, SimpleFSDP di TorchTitan, kernel Pallas specializzati e ottimizzazioni a livello di grafo. TorchTitan offre ricette TorchTPU per il preaddestramento di modelli densi come Llama 3 70B e 405B, Qwen 3 8B e 32B e MuseGlimmer 30B, oltre al modello mixture-of-experts Qwen 3 235B e al trasformatore di diffusione FLUX.1 12B. Google ha dichiarato che l’integrazione è stata convalidata su un massimo di 8.000 TPU Ironwood (TPU v7x) e che l’addestramento TorchTPU tramite TorchTitan ha raggiunto un utilizzo dei FLOP del modello fino al 52% con Qwen 3 8B, con una scalabilità quasi lineare su un massimo di 8.192 TPU Ironwood.

Google ha inoltre integrato nel progetto upstream VeRL, un framework di apprendimento per rinforzo, un backend TPU e una ricetta di riferimento end-to-end. La ricetta GRPO, basata su un’architettura disaggregata trainer-sampler e dimostrata su dataset di ragionamento matematico, è pensata per consentire agli sviluppatori di eseguire i carichi di lavoro VeRL sulle TPU mantenendo il codice esistente dei modelli e degli algoritmi, indipendente dall’hardware.

Ray aggiunge il supporto nativo a TorchTPU per orchestrare l’elaborazione dei dati, l’addestramento, l’apprendimento per rinforzo e la distribuzione di PyTorch su CPU, GPU e TPU all’interno di un unico flusso di lavoro. Edward Oakes, architetto capo di Anyscale, ha dichiarato che le aziende hanno collaborato per fare delle TPU acceleratori di primo livello in Ray, aggiungendo: «TorchTPU integra direttamente le TPU di Google in questo stack, uno standard del settore».

Architettura precedente e disponibilità

TorchTPU è stato presentato pubblicamente il 7 aprile 2026 in un post sul blog di Google Developers, che descriveva un design «Eager First» basato sull’interfaccia PrivateUse1 di PyTorch, con modalità eager Debug, Strict e Fused supportate da una cache di compilazione condivisa. All’epoca, Google ha dichiarato che Fused Eager offre un aumento delle prestazioni dal 50% a oltre il 100% rispetto a Strict Eager e che torch.compile su TorchTPU acquisisce grafi FX con Torch Dynamo e li abbassa tramite XLA e StableHLO, utilizzando kernel personalizzati scritti in Pallas e JAX. Nel post si affermava che lo stack supporta Distributed Data Parallel, FSDPv2 e DTensor fin da subito ed è progettato per gestire l’esecuzione divergente tra i rank. Tra gli elementi della roadmap di TorchTPU per il 2026 era inoltre indicato il lancio di un repository pubblico su GitHub.

Il repository torch_tpu è ora pubblico nell’organizzazione google-pytorch su GitHub ed è distribuito con licenza Apache 2.0. Il README elenca kernel ATen personalizzati, una cache di compilazione per i calcoli XLA, il supporto all’addestramento distribuito con parallelismo dei dati e dei tensori, l’integrazione con torch.compile, utilità di debugging e benchmarking ed esempi per Llama, Qwen, ResNet, DLRM, minGPT e UNet. Il progetto supporta Python 3.11 e versioni successive; il pacchetto torch-tpu 0.2.0 è stato pubblicato su PyPI il 19 ottobre 2026, con pacchetti wheel per Linux x86-64 e CPython dalla versione 3.11 alla 3.14. La guida per gli utenti è disponibile su torchtpu.readthedocs.io e il README specifica che i notebook interattivi del progetto devono essere eseguiti da una macchina virtuale TPU; un’istanza GCE v6e è l’opzione più semplice.

Nello stesso annuncio, Google ha dichiarato che LiteRT, precedentemente noto come TensorFlow Lite, entra a far parte dell’ecosistema PyTorch: la libreria LiteRT-Torch converte un modello PyTorch con una singola litert-torch.convert() chiamata e i modelli convertiti vengono eseguiti su dispositivi Android, iOS, macOS, Windows, Linux, sul Web e su dispositivi IoT. Una suite di agenti TPU, pensata per automatizzare l’onboarding, la diagnostica e l’ottimizzazione delle prestazioni di inferenza per i modelli PyTorch, è disponibile in anteprima iniziale su GitHub.

L’annuncio rimanda gli sviluppatori alla documentazione delle TPU di Google Cloud per iniziare a eseguire PyTorch sulle TPU.