Google et Meta lancent TorchTPU pour exécuter PyTorch en natif sur les TPU
Ce backend open source, lancé chez Google à la mi-2025 avec l'équipe PyTorch de Meta, est déjà utilisé en production par certains des plus grands clients de Google Cloud, selon son équipe.
Illustration IAGoogle et Meta ont présenté TorchTPU, un backend open source qui permet d'exécuter PyTorch en natif sur les puces TPU de Google. L'équipe du projet indique qu'il a vu le jour chez Google à la mi-2025, en collaboration avec l'équipe PyTorch de Meta, et qu'il est en développement actif.
Selon l'équipe, TorchTPU est déjà utilisé en production, pour des charges d'inférence et d'entraînement, par certains des plus grands clients de Google Cloud. L'équipe rappelle que les TPU ont historiquement été liés aux écosystèmes TensorFlow et JAX.
L'équipe présente TorchTPU comme une rupture avec les tentatives précédentes. Google avait publié en 2020 PyTorch/XLA, une première version compilée de PyTorch pour TPU reposant sur les Lazy Tensors. Selon l'équipe, cette intégration exigeait des modifications intrusives du code source, souffrait d'une latence de compilation importante, ne prenait pas en charge le mode d'exécution eager de PyTorch et accusait du retard sur les fonctions de calcul distribué plus récentes. TorchTPU, à l'inverse, fait de l'exécution eager une priorité absolue et n'a aucune dépendance à JAX dans son runtime.
L'objectif affiché est que le passage d'une charge de travail vers un TPU n'exige guère plus que de remplacer, dans le code, l'appareil cible GPU par un TPU. L'équipe prévient qu'un tel changement ne suffirait pas à atteindre les performances maximales du TPU, car les formes, le sharding et les opérations d'un modèle peuvent être optimisés pour les GPU, mais affirme que le modèle devrait tout de même fonctionner prêt à l'emploi.
À ce jour, indique l'équipe, le projet a :
- ajouté la prise en charge des modes eager et compilé, dont un nouveau mode eager appelé DeferAndFuse ;
- ajouté la prise en charge d'environ 450 opérations ATen ;
- activé les noyaux personnalisés écrits avec JAX/Pallas et le DSL Helion ;
- validé l'exécution prête à l'emploi sur les 5 000 modèles Hugging Face les plus téléchargés ;
- validé l'intégration avec vLLM, SGLang, TorchTitan, Ray et TorchRec.
L'équipe énumère également des limites. Un appareil TPU ne peut être accessible qu'à un seul processus en espace utilisateur ; la colocalisation pour l'apprentissage par renforcement ne peut donc pas être prise en charge d'emblée, et l'équipe indique étudier des moyens de contourner ce point. Elle ajoute aussi la prise en charge du dynamisme borné, car le compilateur et le runtime sous-jacents supposent que les formes des tenseurs sont connues à la compilation.
En matière de précision, TorchTPU utilise par défaut un mode intermédiaire qui multiplie en bfloat16 et accumule en FP32. Les modes supérieurs restituent davantage de précision FP32 pour un coût de calcul d'environ 3 fois ou d'environ 6 fois supérieur, et l'équipe recommande le mode par défaut pour la plupart des charges d'apprentissage profond.