Una nuova era per le TPU con PyTorch

Fino a poco tempo fa, chi voleva eseguire PyTorch su una TPU doveva ricorrere a torch_xla o, in alternativa, imparare JAX. Con TorchTPU, un backend nativo per PyTorch, le TPU diventano semplicemente un altro dispositivo, al pari di CPU e GPU. Questo cambiamento apre le porte al supporto completo delle TPU all’interno dell’ecosistema Hugging Face: Transformers, Diffusers e i kernel specializzati funzionano ora in modo nativo.

Basta scegliere un modello dal Hub, impostare il dispositivo a "tpu" e il gioco è fatto. L’intera esperienza di sviluppo rimane quella tipica di PyTorch, senza dover modificare il codice o gestire compilazioni lazy.

Che cosa sono le TPU?

Le Tensor Processing Unit (TPU) sono gli acceleratori AI proprietari di Google, progettati esclusivamente per l’algebra matriciale alla base delle reti neurali. Dopo circa dieci anni di sviluppo, Google le utilizza per addestrare e servire i propri modelli, incluso Gemini.

Le TPU sono disponibili su Google Cloud in configurazioni che vanno da un singolo chip fino a pod composti da migliaia di chip. La loro forza risiede nella capacità di fornire enormi quantità di calcolo, ideale sia per il training su larga scala che per il servizio di modelli in produzione.

Le difficoltà storiche

Il problema principale è sempre stato il software. L’intero stack delle TPU è nato attorno a TensorFlow e JAX, lasciando gli utenti PyTorch in una posizione svantaggiata. Con torch_xla era possibile far funzionare PyTorch, ma il comportamento lazy dei tensori, le ricompilazioni frequenti del grafo e le difficoltà di debug rendevano l’esperienza poco fluida.

Da Optimum TPU a TorchTPU

Per semplificare la vita, Hugging Face aveva introdotto Optimum TPU, che forniva inference e training per i modelli Hugging Face su TPU. Con l’arrivo di TorchTPU, Optimum TPU non è più necessario per far girare Transformers e Diffusers su TPU.

TorchTPU aggiunge le TPU come nuovo backend di PyTorch. Il flusso di lavoro rimane identico a quello tradizionale:

    • Scrivi codice PyTorch standard.
    • Installa TorchTPU con un unico comando pip install torch==2.x+cpu -f https://download.pytorch.org/whl/torch_stable.html (l’indice extra garantisce di scaricare la versione CPU, evitando il pacchetto CUDA non necessario).
    • Utilizza i kernel Pallas ottimizzati per TPU, il supporto al training distribuito e l’integrazione torch.compile per massimizzare le prestazioni.

Transformers su TPU: “ovunque” è reale

La versione 5 di Transformers è stata costruita attorno all’idea di essere la fonte di verità per le definizioni dei modelli, in modo che possano girare ovunque. Ora “ovunque” include anche le TPU. Per caricare un modello su TPU basta passare devicemap="tpu" al metodo frompretrained.

Il modello viene eseguito in modalità eager per impostazione predefinita, ma è possibile attivare torch.compile per ottenere ulteriori miglioramenti di velocità o avviare un training distribuito con pochi parametri aggiuntivi.

Diffusers e TPU: zero fork, zero sorprese

Le TPU sono ora un dispositivo di prima classe anche per Diffusers. Per prototipare o fare debug, basta chiamare pipe.to("tpu"). Per il deployment, si può aggiungere una singola chiamata di compilazione; le pipeline e il codice rimangono invariati, senza la necessità di mantenere un fork specifico per TPU.

Il primo invocazione compila il grafo; le successive, a parità di dimensione dell’immagine, riutilizzano lo stesso grafo, garantendo tempi di latenza costanti.

Modelli di grandi dimensioni e parallelismo tensoriale

Se un modello supera la capacità di un singolo chip, è possibile utilizzare TensorParallelConfig durante il caricamento con frompretrained. Il modello viene suddiviso in parti (sharding) tra più chip, evitando che l’intero modello debba risiedere in un unico posto. Molti pipeline popolari supportano già questa funzionalità; per quelle che non la includono, è sufficiente aggiungere atpplan.

Esempio di configurazione tipica:

    • TP=4 (modello suddiviso su quattro chip)
    • bf16 precisione
    • batch size = 1

Kernels Pallas e Tokamax: potenza a basso livello

Il Kernels Hub mette a disposizione kernel pronti all’uso per TPU. Si caricano tramite la libreria Kernels e offrono un immediato boost di velocità. I kernel sono scritti in Pallas, che consente un controllo a basso livello sull’hardware.

È inoltre possibile integrare kernel dalla libreria Tokamax. Un esempio è il kernel ragged‑dot utilizzato in un modello Mixture‑of‑Experts, che dimostra come sia possibile estendere le capacità di TPU con codice personalizzato.

Prestazioni concrete: decodifica su TPU v6e‑8

Con batch size 64, la decodifica su TPU v6e‑8 è risultata:

    • 21,5× più veloce per Qwen3‑30B‑A3B
    • 12,8× più veloce per Gemma 4‑26B‑A4B

Questi risultati provengono dal confronto tra il kernel del Hub e l’implementazione predefinita, evidenziando il vantaggio di utilizzare kernel ottimizzati.

Qualità e test continui

Il team di Hugging Face sta costantemente migliorando il supporto TPU eseguendo la suite di test di Transformers su TPU. È in programma l’introduzione di una CI notturna dedicata alle TPU, con i risultati pubblicati sul dashboard di Transformers CI.

Al momento, il 96,5% dei test passa su TPU. Ogni fallimento genera uno script di riproduzione e un issue sul repository TorchTPU; finora sono stati segnalati 34 problemi, di cui 19 già risolti. Alcuni bug scoperti su TPU hanno portato a correzioni anche nel core di Transformers, ad esempio nella StaticCache e nei piani tensor‑parallel di Qwen3.

Conclusioni: le TPU non sono più un “side quest”

Le TPU hanno lasciato il ruolo di periferia per gli utenti PyTorch. Installando TorchTPU, scegliendo un modello dal Hub e impostando device="tpu", è possibile sfruttare la potenza di questi acceleratori senza sacrificare la semplicità di PyTorch.

Hugging Face non vede l’ora di vedere le nuove creazioni della community. In caso di problemi, è consigliato aprire un issue sul repository TorchTPU e condividere la propria esperienza.