TorchTPU e PyTorch su TPU: guida all’installazione e alla migrazione nel 2026

Passare da cuda a tpu in PyTorch richiede relativamente poche modifiche con TorchTPU. Il risultato, però, dipende dall’accesso alla TPU, dal pacchetto scelto, dalla compilazione e dalle forme dei tensori.

Cosa offre davvero TorchTPU

TorchTPU è un backend PyTorch nativo per le TPU Google: non è un’API per modelli ospitati e non sostituisce una TPU VM. L’annuncio di PyTorch parla di esecuzione eager, torch.compile, primitive distribuite e supporto per circa 450 operazioni ATen; il repository pubblico elenca kernel ATen, cache della compilazione, esempi distribuiti e integrazione con torch.compile.

Il percorso di migrazione resta volutamente familiare:

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

È un test di portabilità, non una garanzia sulle prestazioni. Come sottolinea l’annuncio di PyTorch, per sfruttare al massimo una TPU contano ancora le forme del modello, lo sharding e le operazioni più adatte all’hardware.

Cosa serve prima dell’installazione

TorchTPU è utile solo in un ambiente collegato a una TPU Google. Il repository consiglia di eseguire i notebook su una TPU VM e indica un’istanza Google Compute Engine v6e come punto di partenza più semplice. Su un portatile privo di TPU si possono installare i pacchetti, ma non verificare l’esecuzione sul dispositivo.

Il README pubblico supporta Python 3.11+ e indica torch_tpu==0.2.0. Conviene usare un ambiente isolato, così l’esperimento non eredita una build CUDA o dipendenze non pertinenti.

Installare torch_tpu senza trascinarsi dietro CUDA

Crea un ambiente virtuale e installa la wheel CPU di PyTorch insieme al pacchetto TorchTPU bloccato alla versione prevista:

python -m venv .venv
source .venv/bin/activate

pip install torch_tpu==0.2.0 \
  --extra-index-url https://download.pytorch.org/whl/cpu

Come specifica il README, se si omette l’indice aggiuntivo si rischia di installare una wheel CUDA di PyTorch, cioè la dipendenza sbagliata per una TPU VM.

Per la prima migrazione è preferibile usare la wheel rilasciata. La compilazione dai sorgenti richiede Bazel o Bazelisk, oltre ai flag documentati dal repository: --config=wheel_common e, quando necessario, --config=no_rbe.

Verificare il backend con un programma PyTorch minimo

Esegui questo codice sulla TPU VM:

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

Se il risultato viene prodotto e il device indicato è tpu, allocazione dei tensori e moltiplicazione matriciale hanno raggiunto il backend TPU. Se il dispositivo non è disponibile, controlla prima il runtime della VM e l’installazione dei pacchetti, invece di intervenire subito sul codice del modello.

TorchTPU supporta l’esecuzione eager, utile per il debugging, e torch.compile per l’esecuzione compilata. Nell’annuncio di PyTorch si legge che, nei test del progetto, la modalità DeferAndFuse può risultare più veloce del 50% rispetto all’eager strict; è però un risultato dichiarato dal progetto, non un benchmark valido in modo universale.

Migrare un workload PyTorch reale

Parti dalla modifica minima che funziona. Solo dopo aver verificato il percorso su un singolo dispositivo aggiungi la configurazione distribuita specifica per TPU.

1. Sostituisci la selezione del dispositivo con torch.device("tpu") oppure usa torch.accelerator.current_accelerator() dove disponibile.
2. Sposta modello e tensori di input sullo stesso dispositivo.
3. Esegui un batch breve prima di attivare compilazione o modalità distribuita.
4. Mantieni stabili batch size e forme dei tensori durante le misurazioni: cambiare forma può innescare nuova attività di compilazione.
5. Verifica il percorso rappresentativo del modello e lo schema distribuito prima di impegnarti in una migrazione più ampia.

Un frammento CUDA distribuito usa spesso NCCL. L’esempio pubblicato per TorchTPU utilizza invece il backend del process group dedicato alle TPU:

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

Il comando esatto per l’avvio distribuito dipende dall’ambiente TPU e dal workload. Cambiare la stringa del dispositivo non configura automaticamente sharding, assegnazione dei processi o pipeline di input.

La prima esecuzione può comportarsi diversamente dalle successive. La cache di compilazione documentata da TorchTPU consente di riutilizzare grafi ripetuti: misura quindi separatamente il warm-up e il throughput a regime.

Dove TorchTPU è pronto e dove deve ancora maturare

Google documenta le integrazioni di TorchTPU con vLLM, SGLang, TorchTitan e gli strumenti Hugging Face. Il suo annuncio sull’inferenza descrive percorsi TPU nativi per vLLM e SGLang, inclusi lo spostamento della KV cache e le funzionalità di serving. L’annuncio Cloud di Google riporta il risultato del 52% di MFU ottenuto da Qwen 3 8B con TorchTitan e riferimenti a utilizzi in produzione.

I limiti pratici sono però meno rifiniti di quanto lascino intendere i titoli degli annunci:

Quando conviene adottare TorchTPU

Vale la pena provare TorchTPU già oggi se hai accesso a una TPU Google, esegui modelli PyTorch e vuoi usare un framework comune per training o inferenza. Il caso d’uso iniziale più solido è quello di un team disposto a profilare e ottimizzare le forme dei tensori, senza aspettarsi che una modifica di una sola riga garantisca un’efficienza equivalente a quella di una GPU.

Meglio rimandare la migrazione di un workload critico se servono una matrice di compatibilità matura, tempi di cold start prevedibili o una funzione da cui dipende lo stack ma che TorchTPU non espone ancora. Procedi per gradi: un modello, una TPU slice, un set di valutazione fisso e misurazioni separate per tempo di installazione, prima compilazione, throughput a regime e correttezza dell’output.

Domande frequenti su TorchTPU

TorchTPU è la stessa cosa di PyTorch/XLA?

No. TorchTPU è un backend PyTorch più recente, progettato attorno all’esecuzione eager e a torch.compile; PyTorch/XLA è il precedente percorso di integrazione con le TPU. Il codice esistente basato su XLA può richiedere una valutazione di migrazione separata.

Posso installare TorchTPU su un portatile senza TPU?

Puoi installare i pacchetti, ma non puoi verificare l’esecuzione su TPU senza un’infrastruttura Google TPU compatibile. La configurazione pratica del repository ufficiale presuppone una TPU VM.

Perché il primo step con TorchTPU è stato lento?

La prima esecuzione può compilare un grafo e riempire le cache. Misura separatamente le iterazioni successive e mantieni stabili le forme dei tensori mentre analizzi eventuali ricompilazioni.

Sostituire cuda con tpu garantisce le stesse prestazioni?

No. La sostituzione del dispositivo serve a verificare la correttezza del porting; prima di confrontare le prestazioni devi profilare forme, sharding, copertura degli operatori e warm-up.

TorchTPU è pronto per la produzione?

Google descrive utilizzi in produzione per inferenza e training, ma componenti e funzionalità dell’ecosistema hanno livelli di maturità diversi. Prima di impegnarti, testa il modello, il motore di serving e lo schema distribuito esatti che intendi usare.

Installa la wheel bloccata alla versione prevista su una TPU VM, esegui la verifica 2×2 e confronta un workload a regime con il percorso GPU attuale. In questo modo distingui la portabilità del codice, il warm-up del compilatore, la copertura degli operatori e l’economia dell’hardware.