Google a publié TorchTPU en open source. Cette couche logicielle permet aux développeurs d’exécuter PyTorch — l’outil largement utilisé pour créer et faire tourner des modèles d’IA — sur les Tensor Processing Units de Google, ou TPU. Développé avec l’équipe PyTorch de Meta, le projet vise à transférer des charges de travail existantes sur du matériel TPU avec relativement peu de changements de code, plutôt qu’à tout réécrire.
Selon Google, TorchTPU prend en charge l’exécution eager, qui lance les opérations au fur et à mesure et facilite le débogage, ainsi que l’exécution compilée pour les charges de travail en production. Les développeurs peuvent remplacer le paramètre d’appareil « cuda » par « tpu » et utiliser le backend distribué TPU ; Google cite aussi l’API de détection des accélérateurs de PyTorch comme solution de remplacement. TorchTPU est disponible sur GitHub, mais sa feuille de route prévoit encore des travaux sur les formes d’entrée dynamiques, la vitesse de compilation, la prise en charge des opérateurs, l’apprentissage par renforcement et la quantification.
Qu’est-ce qui change pour les développeurs ?
Cette sortie rend les TPU compatibles avec des logiciels déjà utilisés par de nombreux développeurs PyTorch. Google indique que TorchTPU fonctionne avec les outils d’inférence vLLM et SGLang, ainsi qu’avec des projets d’entraînement et d’orchestration comme TorchTitan, Ray et TorchRec. Dans Transformers v5, Hugging Face explique que les développeurs peuvent sélectionner un TPU avec le paramètre device_map="tpu". Sa bibliothèque Diffusers prend également en charge l’exécution sur TPU, y compris en mode compilé et avec la répartition du travail entre plusieurs accélérateurs.
Cette approche diffère de PyTorch/XLA, la solution proposée jusque-là par Google pour faire tourner PyTorch sur ses TPU. Google affirme que ce système nécessitait des changements de code plus importants, entraînait de longs temps de compilation et ne prenait pas en charge le mode eager de PyTorch. Google a mis en place un pipeline de tests qui, selon InfoQ, évalue environ 5 000 modèles Hugging Face couvrant plus de 500 architectures. Cela témoigne d’une large campagne de tests, mais ne garantit pas que tous les modèles ou toutes les charges de travail fonctionneront sans modification.
Comment se compare-t-il aux autres options matérielles ?
Google indique que le mode DeferAndFuse de TorchTPU — un hybride entre exécution eager et compilée — peut être au moins 50 % plus rapide que son mode Strict Eager, tout en réduisant l’utilisation maximale de mémoire de l’accélérateur. Lors d’un autre test, Google a exécuté un modèle Muse Glimmer de 30 milliards de paramètres en PyTorch standard : InfoQ rapporte un débit d’inférence de 25 700 tokens par seconde sur une puce TPUv7, avec une longueur de séquence de 5 120, et de 22 400 à 10 240. L’entraînement sur quatre puces TPUv7 a atteint respectivement 5 600 et 4 600 tokens par seconde pour ces longueurs. Il s’agit de résultats communiqués par Google, et non d’une comparaison indépendante.
TechCrunch rapproche TorchTPU de TorchNeuron, annoncé la même semaine par Amazon pour exécuter des charges de travail PyTorch sur ses puces Trainium. Les deux projets visent à faciliter le recours à des solutions autres que les GPU Nvidia. La sortie de TorchTPU offre une option matérielle supplémentaire aux utilisateurs de PyTorch, mais ne démontre pas à elle seule que la migration d’une charge de travail réduira les coûts ou surpassera la configuration actuelle d’un développeur.
Pourquoi c'est important
Pour les développeurs qui utilisent déjà PyTorch, TorchTPU pourrait faciliter le test ou le déploiement de charges de travail sur les TPU de Google sans avoir à tout reconstruire. On ne sait pas encore si ce changement améliorera le coût ou la vitesse pour un modèle donné ; les chiffres de performance cités proviennent de Google.
