Présentation de TorchTPU

TorchTPU ajoute TPU comme nouveau backend à PyTorch, éliminant la dépendance à torch_xla. L’installation se résume à une seule commande : 

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
. Le paramètre --extra-index-url pointe vers la version CPU, évitant le téléchargement du package CUDA, inutile sur les machines TPU.

Le backend s’appuie sur les kernels Pallas, optimisés pour le calcul matriciel des TPU, et propose une intégration native avec torch.compile pour les scénarios où la compilation JIT améliore les performances.

Intégration dans l’écosystème Hugging Face

Les bibliothèques Transformers (v5) et Diffusers reconnaissent désormais le dispositif "tpu". Un modèle se charge simplement avec device_map="tpu" ; aucune modification du code n’est requise. Exemple d’inférence :

from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="tpu",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))

Le même principe s’applique aux pipelines de diffusion : pipe.to("tpu") bascule le modèle en mode eager, tandis que pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False) active la compilation JIT. La première compilation crée le graphe, les appels suivants le réutilisent tant que la taille d’image reste identique.

Gestion du parallélisme et performances

Pour les modèles dépassant la capacité d’une puce, TensorParallelConfig permet le sharding au moment du chargement, répartissant les poids sur plusieurs TPU sans nécessiter de copie complète en mémoire. Les benchmarks fournis par Hugging Face montrent des gains de débit importants : sur un TPU v6e‑8, le décodage de Qwen3‑30B‑A3B atteint 21,5× la vitesse CPU à un batch de 64, et Gemma 4 26B‑A4B réalise 12,8× d’accélération.

Les kernels additionnels, comme le ragged‑dot de la bibliothèque Tokamax, s’intègrent via le module kernels. Cette extensibilité permet d’optimiser des architectures spécialisées (Mixture‑of‑Experts, etc.) sans modifier le code du modèle.

Fiabilité et suivi de la compatibilité

Le test‑suite de Transformers s’exécute désormais sur TPU, avec un taux de réussite de 96,5 % sur les tests concernés. Chaque échec génère un script de reproduction et ouvre un ticket sur le dépôt TorchTPU ; à ce jour, 34 tickets ont été créés, dont 19 résolus. Cette boucle de rétroaction a déjà corrigé des bugs internes à Transformers (ex. : StaticCache, plans tensor‑parallel de Qwen3).

En résumé, TorchTPU transforme les TPU d’un « side quest » en dispositif de première classe pour les utilisateurs PyTorch, en simplifiant le déploiement, en offrant des performances mesurables et en assurant une couverture de test progressive.