Google ha reso open source TorchTPU, un livello software che consente agli sviluppatori di eseguire PyTorch, il toolkit ampiamente utilizzato per creare e gestire modelli di IA, sulle Tensor Processing Unit di Google, o TPU. Sviluppato insieme al team PyTorch di Meta, è pensato per trasferire i carichi di lavoro esistenti sull’hardware TPU con modifiche relativamente limitate al codice, anziché riscriverli.
Google afferma che TorchTPU supporta sia l’esecuzione eager, che esegue le operazioni man mano che vengono chiamate ed è utile per il debugging, sia quella compilata, destinata ai carichi di lavoro in produzione. Gli sviluppatori possono modificare l’impostazione del dispositivo da “cuda” a “tpu” e usare il backend distribuito di TPU; Google indica anche l’API di PyTorch per il rilevamento degli acceleratori come alternativa. TorchTPU è disponibile su GitHub, ma la roadmap prevede ancora interventi sulle forme dinamiche degli input, sulla velocità di compilazione, sulla copertura degli operatori, sul reinforcement learning e sulla quantizzazione.
Cosa cambia per gli sviluppatori?
Il lancio collega le TPU a software già utilizzati da molti sviluppatori PyTorch. Google afferma che TorchTPU è compatibile con gli strumenti di serving vLLM e SGLang e con progetti di training e orchestrazione tra cui TorchTitan, Ray e TorchRec. In Transformers v5, Hugging Face spiega che gli sviluppatori possono selezionare una TPU con l’impostazione device_map="tpu". Anche la sua libreria Diffusers supporta l’esecuzione su TPU, inclusa la modalità compilata e la distribuzione del lavoro su più acceleratori.
È un approccio diverso da PyTorch/XLA, la precedente soluzione di Google per eseguire PyTorch sulle TPU. Secondo Google, quel sistema richiedeva modifiche al codice più invasive, aveva tempi di compilazione lunghi e non supportava la modalità eager di PyTorch. Google ha realizzato una pipeline di test che, secondo InfoQ, valuta circa 5.000 modelli Hugging Face distribuiti su oltre 500 architetture. È un segnale di test su ampia scala, ma non significa che ogni modello o carico di lavoro funzionerà senza modifiche.
Come si confronta con le altre opzioni hardware?
Google riferisce che la modalità DeferAndFuse di TorchTPU, un ibrido tra esecuzione eager e compilata, può essere almeno il 50% più veloce della modalità Strict Eager e ridurre al contempo il picco di memoria usata dall’acceleratore. In un test separato, Google ha eseguito un modello Muse Glimmer da 30 miliardi di parametri con PyTorch standard: InfoQ riporta un throughput di inferenza di 25.700 token al secondo su un singolo chip TPUv7 con una lunghezza di sequenza di 5.120, e di 22.400 con una lunghezza di 10.240. L’addestramento su quattro chip TPUv7 ha raggiunto rispettivamente 5.600 e 4.600 token al secondo. Sono risultati comunicati da Google, non il frutto di un confronto indipendente.
TechCrunch colloca TorchTPU accanto a TorchNeuron di Amazon, annunciato nella stessa settimana per eseguire carichi di lavoro PyTorch sui chip Trainium dell’azienda. Entrambi puntano a rendere più semplice l’uso di alternative alle GPU Nvidia. L’uscita di TorchTPU aggiunge un’altra opzione hardware per chi usa PyTorch, ma non dimostra di per sé che trasferire un carico di lavoro riduca i costi o superi le prestazioni della configurazione attuale di uno sviluppatore.
Perché conta
Per gli sviluppatori che già usano PyTorch, TorchTPU potrebbe semplificare la sperimentazione o la distribuzione dei carichi di lavoro sulle TPU di Google, senza doverli ricostruire da zero. Non è ancora chiaro se il passaggio migliori costi o velocità per un modello specifico; i dati sulle prestazioni citati provengono da Google.
