Google a officiellement publié TorchTPU en tant que logiciel open source, une pile qui permet d’exécuter PyTorch de manière native sur ses TPU, les accélérateurs IA qu’elle propose en location via Google Cloud. Annoncé le 20/10/2026 lors de la conférence PyTorch 2026, ce projet, développé en partenariat avec l’équipe PyTorch de Meta, est disponible sur GitHub sous licence Apache 2.0. Google affirme qu’avec TorchTPU, les TPU deviennent une option matérielle native pour PyTorch, à l’instar des GPU et CPU.

Les essentiels de TorchTPU en 20 secondes

• Google publie TorchTPU en open source sous licence Apache 2.0, développé avec l’équipe PyTorch de Meta.
• Permet d’exécuter PyTorch sur TPU avec un minimum de modifications dans le code, en mode eager ou avec torch.compile.
• Des outils comme vLLM, SGLang, Hugging Face, Ray et llm-d l’ont déjà intégré.
• Google parle d’un scaling presque linéaire jusqu’à 8 192 TPU Ironwood.

L’annonce, signée par Mark Lohmeyer, vice-président et directeur général de l’infrastructure IA et de la computation chez Google Cloud, ainsi que Bill Jia, vice-président en ingénierie, adresse un des principaux freins à l’adoption des TPU : la majorité des applications IA tournent sous PyTorch, alors que le logiciel natif de Google pour ses puces est traditionnellement basé sur JAX et TensorFlow. Pour les utilisateurs de la cloud, la question est de savoir s’il est possible de transférer une charge de travail à un autre accélérateur sans devoir tout réécrire.

Ce qui change par rapport à PyTorch/XLA

Jusqu’à présent, l’option pour utiliser PyTorch avec TPU était PyTorch/XLA, décrit par le blog des développeurs de Google comme le prédécesseur de TorchTPU, prenant en charge uniquement le code SPMD pur (même programme sur tous les nœuds). TorchTPU s’intègre à PyTorch via l’interface PrivateUse1, réservée par le framework pour les backends matériels externes, ce qui signifie que les tenseurs sur TPU sont des tenseurs PyTorch classiques, sans superpositions ni sous-classes.

Google propose un flux de travail en deux étapes : débogage en mode eager, opération par opération, puis passage à torch.compile pour la production. Lors de cette seconde étape, TorchTPU utilise TorchDynamo pour capturer le graphe et le compiler avec XLA, au lieu de TorchInductor, le compilateur par défaut de PyTorch. Selon Google, il existe aussi un mode eager avec fusion automatique des opérations, apportant une amélioration de 50 % à plus de 100 % par rapport au mode eager classique. Pour l’entraînement distribué, il supporte DDP, FSDPv2, DTensor, et permet d’utiliser des kernels custom en Pallas ou JAX.

La version disponible sur PyPI est la 0.2.0, compatible avec Python 3.11 ou supérieur, et s’installe via pip :

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

Scaling sur Ironwood et stack pour l’inférence sur GKE

Les performances revendiquées par Google proviennent de ses propres tests. Avec TorchTitan, l’implémentation de référence de PyTorch pour l’entraînement distribué, la société affirme avoir atteint une utilisation FLOPs de 52 % en entraînant Qwen 3 8B, avec un scaling quasi-linéaire jusqu’à 8 192 TPU Ironwood (TPU v7x), la génération utilisée pour l’inférence à grande échelle présentée par Google. Le torchTitan couvre aussi des modèles tels que Llama 3 de 70 000 à 405 000 millions de paramètres, Qwen 3 235B avec architecture MoE, et FLUX.1 12B.

En inference, Google a développé des backends natifs pour vLLM, en collaboration avec Inferact, et pour SGLang, avec RadixArk. Au-dessus, on trouve llm-d 1.0, un projet en phase sandbox de la CNCF impliquant Red Hat, IBM et plus de 500 collaborateurs de 25 organisations, intégrant un routage intelligent de la mémoire cache de préfixes et la séparation des phases de pré-remplissage et de décodage. Sur GKE, le service Kubernetes de Google Cloud, il est possible de déployer des modèles sur GPU ou TPU à partir d’un point d’accès unique.

Par exemple, Google cite Zenteiq.ai, qui, selon son fondateur, a multiplié par plus de quatre ses performances d’inférence grâce à GKE Inference Gateway et llm-d, passant de 1 470 à 5 970 millions de tokens par jour par chaque 1 000 TPU v6e. Ray, l’orchestrateur d’Anyscale, supporte aussi TorchTPU pour répartir la charge entre CPU, GPU et TPU.

Une alternative à CUDA, avec ses nuances

Ce mouvement influence clairement le marché des accélérateurs. La majorité des avantages de NVIDIA résident dans CUDA, et PyTorch fonctionne généralement mieux sur leurs GPU. Les TPU compétissent désormais en prix et performance, comme évoqué lors de l’analyse de la rivalité entre TPU Google et NVIDIA, mais le logiciel restait une barrière. Meta, qui utilise des TPU Google, souhaite éviter une dépendance à un seul fournisseur. Jana van Greunen, directrice de l’ingénierie PyTorch chez Meta, affirme que la société ne veut pas être limitée à une architecture unique, et que TorchTPU facilite le déploiement de leurs charges PyTorch sur TPU pour les modèles de langage, la publicité et les systèmes de recommandation.

Restent à démontrer plusieurs aspects hors des tests internes de Google. Leur feuille de route inclut la réduction du besoin en recompilations lors de modifications de la longueur de séquence ou du batch, le support complet des formes dynamiques dans torch.compile, et l’extension de la bibliothèque de kernels précompilés. Ces défis, propres aux compilateurs basés sur XLA, impactent notamment l’inférence avec des charges très variables. De plus, TorchTPU fonctionne uniquement avec les TPU de Google, disponibles via Google Cloud, ce qui limite la portabilité entre GPU et TPU dans PyTorch, sans prise en charge inter-cloud.

Questions fréquemment posées

Qu’est-ce que TorchTPU ?

La pile logicielle de Google permettant d’exécuter PyTorch nativement sur ses TPU. Open source, sous licence Apache 2.0, elle a été développée en collaboration avec l’équipe PyTorch de Meta.

TorchTPU remplace-t-il PyTorch/XLA ?

Google décrit PyTorch/XLA comme son prédécesseur. TorchTPU s’intègre en tant que backend natif dans PyTorch, mais Google n’a pas encore officiellement annoncé son abandon.

Quels sont les performances de TorchTPU ?

D’après Google, 52 % d’utilisation de FLOPs lors de l’entraînement de Qwen 3 8B avec TorchTitan, avec un scaling quasi-linéaire jusqu’à 8 192 TPU Ironwood. Ces chiffres sont fournis par la société.

Peut-on utiliser TorchTPU en dehors de Google Cloud ?

Le logiciel est open source, mais fonctionne uniquement sur les TPU de Google, qui se louent via Google Cloud.

Sources :