Si vous avez déjà essayé d'exécuter PyTorch sur un TPU, vous avez probablement porté votre code vers torch_xla, ou vous avez fini par devoir apprendre JAX. Avec TorchTPU, un backend PyTorch natif, le TPU n'est désormais qu'un simple périphérique de plus. Cela ouvre la prise en charge native des TPU dans l'écosystème Hugging Face. Transformers, Diffusers et Kernels s'exécutent nativement sur TPU. Choisissez un modèle sur le Hub, envoyez-le sur "tpu", et c'est tout.
Attendez, c'est quoi un TPU ? 🤔
TPU signifie Tensor Processing Unit. Les TPU sont les accélérateurs d'IA conçus par Google. Ce sont des puces bâties pour une seule tâche : les calculs matriciels sur lesquels les réseaux de neurones passent la majeure partie de leur temps. Google les conçoit depuis une dizaine d'années et les utilise pour entraîner et servir ses propres modèles, dont Gemini.
Les TPU sont disponibles sur Google Cloud, depuis une seule puce jusqu'à des pods de milliers de puces. Ils brillent lorsque vous avez besoin de beaucoup de puissance de calcul, pour l'entraînement ou pour le service à grande échelle.
Le hic a toujours été le logiciel. La pile logicielle des TPU s'est développée autour de TensorFlow et JAX, et les utilisateurs de PyTorch devaient fournir plus d'efforts pour y accéder.
L'ancienne méthode
Jusqu'à présent, exécuter PyTorch sur des TPU signifiait passer par PyTorch/XLA (torch_xla). Cela fonctionnait, mais avec ses singularités. Les tenseurs étaient paresseux (les tenseurs, pas vous), si bien que votre code ne s'exécutait pas toujours quand vous le pensiez. De petits changements pouvaient déclencher des recompilations de graphes, et le débogage était difficile à raisonner.
Pour vous faciliter la vie, nous avons créé Optimum TPU. Il offrait aux modèles Hugging Face la prise en charge de l'inférence et de l'entraînement sur les TPU. Maintenant que TorchTPU apporte une prise en charge native, Optimum TPU n'est plus nécessaire pour exécuter les modèles Hugging Face avec Transformers et Diffusers.
Voici TorchTPU 🚀
TorchTPU ajoute le TPU comme nouveau backend à PyTorch. Vous écrivez du code PyTorch ordinaire. Il est livré avec des kernels Pallas optimisés pour les TPU, la prise en charge de l'entraînement distribué et une intégration torch.compile pour quand vous voulez plus de vitesse. Son installation sur une VM TPU ne nécessite qu'une commande :
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
L'index supplémentaire vous fournit la version CPU de torch. Sans lui, pip récupère la version CUDA, bien plus volumineuse, dont vous n'avez pas besoin sur un TPU.
Vos bibliothèques préférées, désormais sur TPU
Transformers
Transformers v5 a été conçu autour d'une idée : être la source de vérité des définitions de modèles, afin qu'ils puissent s'exécuter partout. « Partout » inclut désormais les TPU. Pour charger un modèle sur un TPU, passez device_map="tpu". C'est tout.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu", # will load model directly on the 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 modèle s'exécute en mode eager par défaut. Vous pouvez aussi utiliser torch.compile pour de meilleures performances, ou l'exécuter dans une configuration distribuée.
Images ? Vidéos ? Diffusers les peint sur les TPU 🎨
TPU est désormais un dispositif de première classe dans Diffusers. Pour le prototypage ou le débogage, le mode eager ne nécessite que pipe.to("tpu"). Pour le déploiement, le mode compilé ajoute un appel compile . Vos pipelines et votre code restent identiques, et il n'y a pas de fork TPU à maintenir.
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu") # eager mode: that's the whole change
# optional: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]
Le premier appel compile, et chaque appel suivant avec la même taille d'image réutilise le graphe.
Des modèles trop gros pour une seule puce ? Passez un TensorParallelConfig à from_pretrained et le transformer est shardé à travers les puces pendant le chargement, si bien que le modèle complet n'a jamais besoin de tenir au même endroit. De nombreux pipelines populaires le prennent en charge dès l'origine, et ceux qui ne le font pas sont à seulement un tp_plan de distance. Le guide TPU de Diffusers contient la recette complète.
Voici quelques exemples exécutés sur un Cloud TPU v6e-4 :
TP=4 (chaque modèle réparti sur quatre puces) · bf16 · batch 1
| MiniMax-H3 (vidéo + audio) | Qwen-Image-2.1 | |
|---|---|---|
| « La mascotte en peluche duveteuse de Hugging Face issue de l'image, un visage jaune rond et souriant avec des joues rosées et deux petites mains ouvertes, se tient à la proue d'une trière grecque antique […] Bruit des vagues qui se brisent, craquement du bois, vent dans les gréements, mouettes au loin, et une montée orchestrale épique et grave. » · 1344×768, 5.2 s, 49 étapes · 188 s compilé, 262 s eager | « Un capybara portant un chapeau de sorcier, lisant un livre à la lueur d'une bougie, peinture à l'huile » · 1024×1024, 40 étapes · 3.5 s compilé, 8.0 s eager | « Une enseigne de magasin au néon qui affiche 'QWEN IMAGE 2.1', nuit pluvieuse, reflets sur le trottoir mouillé » · 2048×2048, 40 étapes |
Kernels
Le Kernels Hub propose des kernels TPU prêts à l'emploi. Chargez-en un avec la bibliothèque Kernels et offrez un gain de vitesse à votre modèle.
Les noyaux TPU sont écrits en Pallas, ce qui vous donne un contrôle de bas niveau sur le matériel. Vous pouvez également importer des noyaux depuis la bibliothèque Tokamax . Voici un noyau ragged-dot de Tokamax intégré à un modèle Mixture-of-Experts :
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-30B-A3B-Instruct-2507",
experts_implementation="tokamax_ragged_dot",
)
Plus le batch est grand, plus le gain est important. Avec une taille de batch de 64, le décodage est 21,5× plus rapide sur Qwen3-30B-A3B et 12,8× plus rapide sur Gemma 4 26B-A4B.
Débit de décodage sur un TPU v6e-8, noyau Hub par rapport à l'implémentation par défaut.
Nous testons pour que vous n'ayez pas à le faire 🧪
Nous améliorons continuellement la prise en charge des TPU dans transformers en exécutant la suite de tests de la bibliothèque. Nous travaillons à l'ajout d'une CI nocturne sur TPU et à la mise à disposition des résultats dans le (tableau de bord de transformers)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].
Le tableau de bord CI de Transformers, avec les résultats TPU à côté des autres accélérateurs. Travail en cours !
Ce n'est pas encore tout vert, et cela ne nous dérange pas. Parmi les tests exécutés sur TPU, 96,5 % passent. Chaque échec que nous trouvons devient un petit script de reproduction et une issue sur le dépôt TorchTPU. À ce jour, nous avons signalé 34 issues, et 19 sont déjà corrigées. L'exécution sur TPU a également permis de détecter quelques bugs dans Transformers lui-même, comme dans StaticCache et les plans tensor-parallèles de Qwen3. Nous les avons corrigés en amont.
Si quelque chose se casse sur TPU, nous voulons le découvrir avant vous.
Commencer
Les TPU ne sont plus une quête annexe pour les utilisateurs de PyTorch. Installez TorchTPU, récupérez un modèle depuis le Hub, et envoyez-le vers "tpu".
Nous avons hâte de voir ce que vous construirez 🤗. Si vous rencontrez un problème, ouvrez une issue et parlez-nous-en.