Google ouvre TorchTPU : PyTorch peut désormais exploiter nativement les accélérateurs IA de Google

Google publie TorchTPU en open source pour exécuter PyTorch sur ses TPU, avec le support de vLLM, SGLang, Hugging Face et TorchTitan.

Les développeurs d'intelligence artificielle disposent désormais d'une nouvelle manière d'exploiter les accélérateurs de Google sans réécrire leurs modèles pour un environnement logiciel particulier. À l'occasion de la PyTorch Conference 2026, Google Cloud a publié TorchTPU en open source, une infrastructure permettant d'exécuter nativement des programmes PyTorch sur ses Tensor Processing Units (TPU). Développée avec l'équipe PyTorch de Meta, elle associe les API habituelles du framework à des outils de compilation, d'entraînement distribué et d'inférence. L'annonce du 20 octobre s'accompagne d'intégrations avec Hugging Face, vLLM, SGLang, Ray et TorchTitan. L'enjeu dépasse la compatibilité technique : Google veut rendre ses accélérateurs plus accessibles aux équipes qui construisent déjà leurs modèles autour de PyTorch, tout en leur permettant de choisir entre CPU, GPU et TPU selon leurs besoins. Réduire la dépendance à une architecture matérielle

PyTorch s'est imposé comme l'un des principaux environnements de développement pour les modèles d'apprentissage automatique. Il permet aux chercheurs et aux ingénieurs de définir des architectures, d'organiser leur entraînement et de contrôler les calculs à l'aide d'une interface Python. Mais exploiter efficacement le même programme sur différentes familles d'accélérateurs reste un travail technique important, notamment lorsque les outils de compilation et les bibliothèques spécialisées ne sont pas identiques.

Les TPU de Google disposent depuis plusieurs années d'écosystèmes logiciels permettant d'entraîner et de servir des modèles, notamment à travers JAX, TensorFlow et PyTorch/XLA. TorchTPU ne représente donc pas la première possibilité d'exécuter PyTorch sur un TPU. La nouveauté concerne l'intégration native : le framework propose un backend PyTorch conçu pour utiliser les interfaces existantes avec moins d'adaptations propres au matériel.

Le développeur peut ainsi travailler en mode eager, qui exécute les opérations de manière interactive et facilite le débogage, puis utiliser `torch.compile` pour optimiser les calculs destinés à la production. Les fonctions distribuées permettent également de répartir les tâches entre plusieurs accélérateurs. Cette organisation cherche à rapprocher l'expérience de développement sur TPU de celle déjà connue sur GPU ou CPU.

La démarche bénéficie du soutien de Meta, qui a participé au développement de TorchTPU. Pour les entreprises exploitant de nombreux modèles, l'objectif est de ne pas lier systématiquement leurs applications à un seul fournisseur de processeurs ou à une chaîne logicielle particulière. Cette portabilité reste relative aux fonctionnalités et aux modèles pris en charge, mais elle réduit la nécessité de maintenir plusieurs implémentations d'un même projet. Hugging Face intègre les TPU à Transformers et Diffusers

Le rapprochement avec Hugging Face constitue l'une des principales conséquences pratiques de l'ouverture de TorchTPU. Dans Transformers v5, un modèle compatible peut être chargé sur les accélérateurs de Google en utilisant simplement le paramètre `devicemap="tpu"`. L'utilisateur conserve ainsi une interface familière sans devoir transformer manuellement son projet en application spécialisée pour les TPU.

La bibliothèque Diffusers bénéficie également d'une prise en charge native. Elle peut utiliser le mode eager, la compilation PyTorch et le parallélisme tensoriel pour répartir les modèles génératifs sur plusieurs accélérateurs. Cette intégration concerne notamment les workflows de diffusion utilisés dans la génération d'images et d'autres contenus multimodaux.

Google indique avoir développé des pipelines de tests automatisés exécutés chaque nuit pour vérifier la compatibilité de la suite Transformers avec les TPU. Cette maintenance est importante pour un environnement dans lequel les architectures de modèles et les bibliothèques évoluent rapidement. Elle ne signifie pas pour autant que tous les checkpoints publiés sur Hugging Face fonctionneront immédiatement, sans modification ni optimisation supplémentaire. vLLM et SGLang arrivent sur un backend PyTorch natif

L'inférence, c'est-à-dire l'exécution des modèles pour produire des réponses, constitue l'autre grand volet de cette publication. Google a travaillé avec Inferact sur vLLM-TorchTPU, une déclinaison du moteur d'inférence vLLM destinée aux accélérateurs de Google. L'intégration conserve une grande partie de l'organisation existante du logiciel, notamment ses ordonnanceurs, ses serveurs API et ses mécanismes de gestion des requêtes.

Le backend prend en charge le traitement continu des lots, la quantification et la répartition des modèles sur plusieurs puces. Des kernels développés avec Pallas optimisent certaines opérations de calcul sur TPU. Pour les équipes qui utilisent déjà vLLM, l'intérêt réside dans la possibilité de conserver leurs définitions de modèles et une large partie de leurs outils tout en changeant d'infrastructure matérielle.

Une démarche comparable a été menée avec RadixArk autour de SGLang-TorchTPU. Le moteur est particulièrement utilisé pour les applications agentiques et les générations structurées, où plusieurs requêtes peuvent partager de longues portions de contexte. SGLang exploite notamment RadixAttention pour réutiliser des informations déjà calculées et éviter de traiter plusieurs fois les mêmes préfixes.

Les nouveaux backends utilisent PyTorch directement sur TPU, contrairement aux intégrations antérieures de ces moteurs reposant sur JAX. Selon Google, cette approche facilite le débogage et la personnalisation des environnements d'inférence. Leur publication ne garantit toutefois pas des performances identiques pour tous les modèles ni une parité immédiate avec l'ensemble des fonctions disponibles sur GPU. llm-d organise les requêtes entre TPU et GPU

L'exécution d'un modèle sur plusieurs accélérateurs ne suffit pas à construire un service capable de répondre efficacement à de nombreux utilisateurs. À grande