Hugging Face intègre les TPU de Google à PyTorch

Hugging Face annonce TorchTPU, un backend PyTorch natif pour les unités de traitement Tensor de Google. L’annonce a été publiée le 20 octobre 2026. Cette intégration fait du TPU un périphérique supplémentaire dans PyTorch. Les utilisateurs peuvent charger un modèle depuis le Hub, puis l’envoyer vers « tpu ». Transformers, Diffusers et Kernels disposent désormais d’un support natif sur ces accélérateurs. TorchTPU inclut aussi des noyaux Pallas réglés pour les TPU, l’entraînement distribué et une intégration de torch.compile. Hugging Face présente cette évolution comme une simplification du parcours PyTorch sur TPU. L’installation sur une machine virtuelle TPU utilise une seule commande pip install torch-tpu.

Le logiciel change Les TPU sont les accélérateurs conçus par Google pour les calculs matriciels des réseaux neuronaux. Google les utilise notamment pour entraîner et servir ses modèles, dont Gemini. Ils sont proposés sur Google Cloud, d’une puce unique jusqu’à des ensembles de plusieurs milliers de puces. Avant TorchTPU, les utilisateurs PyTorch passaient par PyTorch/XLA, aussi appelé torch_xla. Cette approche reposait sur des tenseurs différés et pouvait déclencher des recompilations après de petites modifications. Hugging Face avait créé Optimum TPU pour apporter l’inférence et l’entraînement de ses modèles sur TPU. Avec le backend natif, Optimum TPU n’est plus nécessaire pour exécuter Transformers et Diffusers.

TorchTPU ajoute un périphérique PyTorch

Transformers adopte le TPU comme périphérique avec la version 5, conçue pour faire fonctionner les définitions de modèles sur différents environnements. Le chargement d’un modèle nécessite le paramètre device_map="tpu". L’exemple fourni utilise [Qwen](/signal-ia/actu/google-ouvre-torchtpu-et-simplifie-pytorch-sur-ses-tpu)/Qwen3.5-9B, puis transfère les entrées vers ce même périphérique. Le modèle démarre en mode immédiat par défaut. Les utilisateurs peuvent ensuite activer torch.compile ou une exécution distribuée. Diffusers suit le même principe avec pipe.to("tpu"). Un mode compilé ajoute un appel de compilation. Le premier appel compile le graphe, puis les appels suivants réutilisent ce graphe si la taille d’image reste identique.

Les modèles se distribuent aussi sur plusieurs puces grâce à TensorParallelConfig. Cette option fragmente le transformateur pendant le chargement, afin que le modèle complet ne doive pas tenir sur une seule puce. Plusieurs pipelines Diffusers la prennent en charge directement. Hugging Face cite des essais réalisés sur un Cloud TPU v6e-4, avec un parallélisme tensoriel de quatre et un format bfloat16. Sur Qwen-Image-2.1, une image 1024×1024 en 40 étapes prend 3,5 secondes en mode compilé, contre 8 secondes en mode immédiat. La documentation Diffusers pour TPU fournit la procédure complète.

Hugging Face mesure la compatibilité TPU

Les noyaux accélèrent également certains modèles. Le Hub Kernels propose des noyaux TPU écrits en Pallas, avec un contrôle précis du matériel. Les utilisateurs peuvent aussi intégrer des noyaux de Tokamax dans des modèles à mixture d’experts. Sur un TPU v6e-8, le noyau Hub testé accélère le décodage de Qwen3-30B-A3B de 21,5 fois avec une taille de lot de 64. Le gain atteint 12,8 fois sur Gemma 4 26B-A4B dans la même configuration. Hugging Face publie TorchTPU sur GitHub. Google documente également les TPU sur Cloud.

Les tests progressent dans la suite Transformers exécutée sur TPU. Parmi les tests qui s’exécutent sur ce matériel, 96,5 % passent actuellement. Hugging Face a signalé 34 problèmes dans le dépôt TorchTPU, dont 19 sont déjà corrigés. Les essais ont aussi révélé des défauts dans Transformers, notamment dans StaticCache et les plans de parallélisme tensoriel de Qwen3. Ces défauts ont été corrigés en amont. L’équipe travaille enfin à ajouter une intégration continue quotidienne sur TPU. Les résultats doivent apparaître dans le tableau de bord Transformers. Les utilisateurs peuvent installer TorchTPU, choisir un modèle du Hub et le transférer vers « tpu ».