Google ha annunciato TorchTPU il 20 ottobre 2026: un backend open source che permette di eseguire PyTorch sulle sue TPU, sviluppato insieme al team PyTorch di Meta. Per una prova di base bastano l’installazione del pacchetto e la selezione del dispositivo TPU; i job distribuiti richiedono anche impostazioni specifiche per la topologia e la rete della TPU.
La migrazione di base parte dal dispositivo
La documentazione ufficiale di Google Cloud indica di installare TorchTPU con pip install torch_tpu e di selezionare la TPU tramite torch.device("tpu"). L’esecuzione eager, in cui PyTorch esegue le operazioni man mano, è quella predefinita; torch.compile permette invece di compilare il modello.
Nel codice che già sposta il modello su CUDA, il cambiamento essenziale può essere model.to("cuda") → model.to("tpu"). Il team TorchTPU avverte però che modificare il dispositivo non basta necessariamente a ottenere le prestazioni massime: forme dei dati, partizionamento e operazioni possono essere stati ottimizzati per una GPU e richiedere adattamenti sulla TPU.
I job distribuiti richiedono impostazioni TPU
Quando il carico usa più dispositivi, TorchTPU adotta un processo per dispositivo e il backend distribuito tpu_dist. Il lancio deve conoscere la topologia fisica dello slice TPU e gli indirizzi necessari a costruire la rete tra i chip. In Google Cloud si può usare anche torch.tpu.distributed.set_tpu_launch_env() per impostare queste variabili.
| Tipo di esecuzione | Dispositivo o backend | Configurazione di lancio |
| Base | torch.device("tpu") oppure .to("tpu") | Installare torch_tpu; esecuzione eager oppure torch.compile |
| Distribuita | tpu_dist, con un processo per dispositivo | Impostare topologia TPU e indirizzi dello slice builder, oppure usare l’helper documentato per Google Cloud |
Per chi porta un’applicazione già distribuita, quindi, la modifica non si esaurisce nel sostituire nccl con tpu_dist: occorre anche fornire al runtime le informazioni della specifica configurazione TPU.
Che cosa misurano i test pubblicati
Il progetto PyTorch riferisce il supporto di circa 450 operatori ATen, escludendo dal conteggio quelli con delegati strutturali o compositi. Nei test sui modelli, 4.821 su 4.917 del gruppo Tier 1 hanno superato una prova di avvio senza crash, pari al 98%. Nel gruppo Tier 2, 638 dei 657 modelli rappresentativi hanno raggiunto la parità numerica rispetto a una baseline GPU, il 97%.
Per i test sul bounded dynamism, cioè forme dei tensori variabili entro limiti prestabiliti, il progetto ha riportato risultati corretti per 834 degli 839 LLM causali valutati, il 99%. Il team descrive quest’area come ancora in sviluppo; nella situazione illustrata dal progetto, torch.compile imposta dynamic=False sulla TPU.
Un esempio concreto mostra anche dove può servire intervenire: nel porting di Muse Glimmer 30B da GPU, il progetto ha sostituito FlexAttention, non supportata in quel test, con F.scaled_dot_product_attention. È una modifica riferita a quell’implementazione, non una regola per ogni modello.
Il progetto ha inoltre riferito un test di training di Qwen3 MoE 235B su configurazioni fino a 8.192 chip TPU v7. Nel run più grande, con sequenza di 4.096 token e batch di 262.144, ha misurato 2.104 token al secondo per TensorCore, un MFU del 27,1% e 31,2 secondi per passo. Sono risultati di una configurazione specifica del progetto PyTorch, non una misura comparativa tra lo stesso carico su TPU e GPU.
Riutilizzare i pesi salvati
Se il modello PyTorch è compatibile con TorchTPU, puoi ricrearne la definizione, caricare i pesi salvati nel suo state_dict e spostarlo sul dispositivo TPU con .to("tpu"). PyTorch descrive il state_dict del modello come un dizionario di parametri e buffer; per riprendere l’addestramento, lo stato dell’optimizer è un elemento separato e va salvato e caricato a sua volta. Riutilizzare i pesi del modello, quindi, non equivale a ripristinare da solo ogni componente del training.
