Google rend TorchTPU open source et apporte PyTorch en mode natif sur ses TPU

Google a annoncé le 20 octobre 2026, lors de la PyTorch Conference 2026, que TorchTPU, sa pile native PyTorch pour les unités de traitement tensoriel (TPU), était désormais disponible en open source sur GitHub, a indiqué l’entreprise sur le blog Google Cloud. Le logiciel a été développé conjointement avec l’équipe PyTorch de Meta et est conçu pour exécuter des charges de travail PyTorch existantes sur des TPU en n’apportant qu’un minimum de modifications au code.

Avec cette publication, les TPU deviennent, aux côtés des GPU et des CPU, une option native pour PyTorch, a indiqué Google. La prise en charge couvre l’entraînement, le service et le déploiement au moyen d’outils tels que vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face et LiteRT. Les développeurs peuvent déboguer leurs modèles en mode eager, puis passer à des performances de niveau production grâce à torch.compile, afin d’optimiser l’utilisation des TPU et de réduire la surcharge mémoire.

Jana van Greunen, directrice de l’ingénierie PyTorch chez Meta, a déclaré dans l’annonce que TorchTPU permet à Meta de développer des solutions matérielles conjointement et avec souplesse avec son équipe centrale de systèmes d’apprentissage automatique, tout en prenant en charge des charges de travail majeures liées aux LLM, à la publicité et aux systèmes de recommandation. « Cette solution nous permet de déployer nos charges de travail PyTorch sur les TPU de Google, de transmettre directement nos commentaires aux équipes d’ingénierie et de bâtir la flexibilité de l’écosystème dont tout le monde a besoin », a-t-elle déclaré.

Intégrations Hugging et de mise en service

Google s’est associé à Hugging Face pour que les modèles à poids ouverts disponibles sur le Hub Hugging Face s’exécutent nativement avec TorchTPU. Dans Transformers v5, le déploiement d’un modèle sur des TPU s’effectue en transmettant device_map="tpu", et la bibliothèque Diffusers s’exécute nativement sur les TPU, avec prise en charge du mode eager, de torch.compile et du parallélisme tensoriel sur plusieurs accélérateurs. Google a également indiqué avoir mis en place des pipelines CI/CD dédiés aux tests nocturnes continus de la suite Hugging Face Transformers sur une infrastructure TPU. Jeff Boudier, vice-président produit chez Hugging Face, a déclaré que cette collaboration concrétise l’engagement de la plateforme en faveur du choix du matériel.

Pour la mise en service, le backend vLLM natif PyTorch de TorchTPU a été développé en partenariat avec Inferact. Google a indiqué qu’il prend en charge le traitement continu par lots, la quantification et le parallélisme multi-puces sur du matériel TPU à l’aide de noyaux Pallas personnalisés, tout en conservant les modèles de workers, les ordonnanceurs et les serveurs API en amont, de sorte que les piles d’inférence existantes ne nécessitent aucune modification. Simon Mo, PDG et cofondateur d’Inferact, a déclaré que cette intégration apporte aux TPU une prise en charge de premier ordre de vLLM.

Un backend natif PyTorch pour SGLang, développé avec RadixArk, cible les charges de travail structurées et agentiques. Selon l’annonce, SGLang utilise RadixAttention pour réutiliser le cache KV entre les requêtes, afin d’éviter de recalculer les préfixes de requête répétés. La version TorchTPU utilise un backend de transfert désagrégé qui sert également de backend d’E/S pour HiCache. Banghua Zhu, directeur technique et cofondateur de RadixArk, a déclaré que TorchTPU avait accéléré le travail de l’entreprise alors qu’elle déployait la prise en charge de SGLang sur différentes plateformes matérielles.

Google publie également llm-d 1.0 avec Red Hat, IBM et une communauté réunissant plus de 500 contributeurs et 25 organisations, dans le cadre d’un projet Sandbox de la Cloud Native Computing Foundation. L’annonce présente llm-d comme un lien entre l’écosystème de mise en service de PyTorch et le déploiement en production basé sur Kubernetes : son routeur utilise un routage tenant compte du cache des préfixes pour envoyer chaque requête à la réplique qui possède déjà son contexte, et sépare la phase de préremplissage, gourmande en calcul, de la phase de décodage, gourmande en mémoire. La version 1.0 ajoute le contrôle de flux multi-tenant, la haute disponibilité du routeur, des API v1 stables et un routage tenant compte des sessions pour les agents. Sur Google Kubernetes Engine, un point de terminaison llm-d commun peut fonctionner aussi bien sur des GPU que sur des TPU.

Le professeur Sashikumaar Ganesan, fondateur de Zenteiq, a déclaré dans l’annonce que GKE Inference Gateway, optimisé par llm-d, avait fourni un débit d’inférence agrégé plus de 4× supérieur pour la génération de données synthétiques de l’entreprise, faisant passer la capacité quotidienne de génération de 1,47 milliard à 5,97 milliards de jetons pour 1 000 accélérateurs TPU v6e.

Entraînement et orchestration

TorchTitan, l’implémentation de référence native PyTorch pour l’entraînement distribué, comprend désormais une intégration TorchTPU de qualité production, selon Google, fondée sur torch.compile, SimpleFSDP de TorchTitan, des noyaux Pallas spécialisés et des optimisations au niveau du graphe. TorchTitan propose des recettes TorchTPU pour le préentraînement de modèles denses tels que Llama 3 70B et 405B, Qwen 3 8B et 32B, et MuseGlimmer 30B, ainsi que du modèle à mélange d’experts Qwen 3 235B et du transformeur de diffusion FLUX.1 12B. Google a indiqué que l’intégration avait été validée sur un maximum de 8 000 TPU Ironwood (TPU v7x) et que l’entraînement TorchTPU via TorchTitan avait atteint jusqu’à 52 % d’utilisation des capacités de calcul du modèle sur Qwen 3 8B, avec une mise à l’échelle quasi linéaire sur un maximum de 8 192 TPU Ironwood.

Google a également intégré en amont un backend TPU et une recette de référence de bout en bout à VeRL, un framework d’apprentissage par renforcement. La recette GRPO, fondée sur une architecture désagrégée séparant l’entraîneur de l’échantillonneur et présentée sur des jeux de données de raisonnement mathématique, vise à permettre aux développeurs d’exécuter des charges de travail VeRL sur des TPU tout en conservant leur code de modèle et d’algorithme existant, indépendant du matériel.

Ray ajoute la prise en charge native de TorchTPU pour orchestrer le traitement des données, l’entraînement, l’apprentissage par renforcement et le service de modèles PyTorch sur des processeurs, des GPU et des TPU au sein d’un même flux de travail. Edward Oakes, architecte en chef chez Anyscale, a déclaré que les entreprises avaient travaillé à faire des TPU des accélérateurs de premier plan dans Ray, ajoutant : « TorchTPU intègre directement les TPU de Google à cette pile standard du secteur. »

Architecture antérieure et disponibilité

TorchTPU a été présenté au public le 7 avril 2026, dans un article du blog des développeurs de Google décrivant une conception « Eager First » fondée sur l’interface PrivateUse1 de PyTorch, avec les modes d’exécution eager Debug, Strict et Fused reposant sur un cache de compilation partagé. Google indiquait alors que Fused Eager offrait un gain de performances de 50 % à plus de 100 % par rapport à Strict Eager, et que torch.compile sur TorchTPU capturait des graphes FX avec Torch Dynamo avant de les abaisser via XLA et StableHLO, à l’aide de noyaux personnalisés écrits en Pallas et JAX. L’article précisait que la pile prenait en charge d’emblée Distributed Data Parallel, FSDPv2 et DTensor, et qu’elle était conçue pour gérer les divergences d’exécution entre les rangs. Il indiquait également que le lancement d’un dépôt public sur GitHub figurait parmi les objectifs de la feuille de route de TorchTPU pour 2026.

Le dépôt torch_tpu est désormais public au sein de l’organisation GitHub google-pytorch et est distribué sous licence Apache 2.0. Son fichier README répertorie des noyaux ATen personnalisés, un cache de compilation pour les calculs XLA, la prise en charge de l’entraînement distribué avec parallélisme des données et des tenseurs, l’intégration de torch.compile, des utilitaires de débogage et d’évaluation des performances, ainsi que des exemples pour Llama, Qwen, ResNet, DLRM, minGPT et UNet. Le projet prend en charge Python 3.11 et les versions ultérieures, et le paquet torch-tpu 0.2.0 a été publié sur PyPI le 19 octobre 2026, avec des paquets wheel Linux x86-64 pour CPython 3.11 à 3.14. Un guide utilisateur est publié à l’adresse torchtpu.readthedocs.io, et le fichier README précise que les notebooks interactifs du projet doivent être exécutés depuis une VM TPU, une instance GCE v6e étant l’option la plus simple.

Dans la même annonce, Google a indiqué que LiteRT, anciennement TensorFlow Lite, rejoignait l’écosystème PyTorch : la bibliothèque LiteRT-Torch convertit un modèle PyTorch en un seul litert-torch.convert() appel, et les modèles convertis s’exécutent sur Android, iOS, macOS, Windows, Linux, le Web et les appareils de l’Internet des objets. Une suite d’agents TPU destinée à automatiser l’intégration, les diagnostics et l’optimisation des performances d’inférence des modèles PyTorch est en préversion initiale sur GitHub.

L’annonce invite les développeurs à consulter la documentation des TPU de Google Cloud pour commencer à exécuter PyTorch sur des TPU.