Google e Meta rilasciano TorchTPU per eseguire PyTorch in modo nativo sulle TPU
Il backend open source, avviato in Google a metà 2025 con il team PyTorch di Meta, è già usato in produzione da alcuni dei più grandi clienti di Google Cloud, ha dichiarato il suo team.
Illustrazione IAGoogle e Meta hanno presentato TorchTPU, un backend open source che consente a PyTorch di funzionare in modo nativo sui chip TPU di Google. Il team del progetto ha spiegato che è nato in Google a metà 2025 in collaborazione con il team PyTorch di Meta e che è in fase di sviluppo attivo.
Secondo il team, TorchTPU è già usato in produzione per carichi di lavoro di inferenza e addestramento da alcuni dei più grandi clienti di Google Cloud. Il team ha osservato che le TPU sono state storicamente legate agli ecosistemi TensorFlow e JAX.
Il team ha descritto TorchTPU come una rottura rispetto ai tentativi precedenti. Nel 2020 Google rilasciò PyTorch/XLA, una prima versione compilata di PyTorch per TPU basata su Lazy Tensors. Il team ha spiegato che quell'integrazione richiedeva modifiche invasive al codice sorgente, soffriva di una forte latenza in fase di compilazione, non supportava la modalità di esecuzione eager di PyTorch e restava indietro sulle funzionalità di calcolo distribuito più recenti. TorchTPU, invece, tratta l'esecuzione eager come una priorità assoluta e non ha dipendenze da JAX nel proprio runtime.
L'obiettivo dichiarato è che spostare un carico di lavoro su una TPU richieda poco più che cambiare nel codice il dispositivo di destinazione da una GPU a una TPU. Il team ha avvertito che questa modifica da sola non garantirebbe le massime prestazioni della TPU, perché shape, sharding e operazioni di un modello potrebbero essere ottimizzati per le GPU, ma ha detto che il modello dovrebbe comunque funzionare subito, senza modifiche.
Finora, ha spiegato il team, il progetto ha:
- aggiunto il supporto per le modalità eager e compilata, compresa una nuova modalità eager chiamata DeferAndFuse;
- aggiunto il supporto per circa 450 operazioni ATen;
- abilitato kernel personalizzati scritti tramite JAX/Pallas e il DSL Helion;
- convalidato l'esecuzione immediata, senza modifiche, sui 5,000 modelli Hugging Face più scaricati;
- convalidato l'integrazione con vLLM, SGLang, TorchTitan, Ray e TorchRec.
Il team ha elencato anche dei limiti. Un dispositivo TPU può essere usato da un solo processo user-space, quindi la colocation per il reinforcement learning non può essere supportata subito e il team ha detto di essere al lavoro per aggirare il problema. Sta inoltre aggiungendo il supporto per il dinamismo limitato, perché il compilatore e il runtime sottostanti presuppongono che le shape dei tensori siano note in fase di compilazione.
Per quanto riguarda la precisione, TorchTPU usa per impostazione predefinita una modalità media che moltiplica in bfloat16 e accumula in FP32. Le modalità superiori recuperano più precisione FP32 a un costo di calcolo di circa 3 o circa 6 volte superiore, e il team consiglia l'impostazione predefinita per la maggior parte dei carichi di lavoro di deep learning.