Le projet TorchTPU, fruit d’une collaboration entre Google et Meta, offre un support natif pour PyTorch sur les unités de traitement tensoriel (TPU) de Google. Ce travail permet aux développeurs d’exécuter des charges de travail PyTorch sur ce matériel avec des modifications de code minimes, en prenant en charge à la fois l’exécution immédiate et compilée, ainsi que le calcul distribué.

Historiquement, l’approche PyTorch/XLA, mise en place par Google dès 2020 en s’appuyant sur JAX et XLA, présentait des lacunes. Selon InfoQ, cette ancienne méthode exigeait des modifications intrusives du code source pour définir des zones compilables, souffrait d’une latence importante lors de la compilation, et ne supportait pas le mode d’exécution impératif (« eager ») propre à PyTorch.

TorchTPU change cette donne en priorisant l’exécution immédiate, permettant aux applications PyTorch existantes de fonctionner avec peu de réécriture. Sur le plan pratique, migrer d’un GPU à un TPU se résume, dans une certaine mesure, à remplacer la référence de dispositif cuda par tpu et le backend distribué nccl par tpu_dist. Les développeurs peuvent aussi utiliser la nouvelle API torch.accelerator.current_accelerator pour identifier l’accélérateur disponible.

Le projet propose trois modes d’exécution immédiate : strict, debug, et DeferAndFuse. Ce dernier mode, selon Google, est un hybride entre exécution immédiate et compilée, et il est affirmé qu’il peut être « 50 % ou plus plus rapide que Strict Eager tout en réduisant la mémoire HBM de pointe » (InfoQ).

Pour valider cette capacité, Google a testé environ 5 000 modèles Hugging Face couvrant plus de 500 architectures sur des TPU à partir de la génération v5e. Concernant les performances, l’équipe a évalué un modèle Muse Glimmer de 30 milliards de paramètres en PyTorch standard. Sur une seule puce TPUv7, le débit d’inférence a atteint 25 700 tokens/seconde pour une séquence de 5 120, et 22 400 tokens/seconde pour 10 240 tokens. Pour l’entraînement, en utilisant quatre TPUv7, le débit était de 5 600 tokens/seconde et 4 600 tokens/seconde pour les mêmes longueurs de séquence.

Le socle technique de ces outils s’appuie sur des compilateurs comme XLA. Ce compilateur, open-source, prend des modèles de frameworks populaires tels que PyTorch, TensorFlow et JAX, et les optimise pour une exécution performante sur CPU, GPU ou accélérateurs spécialisés, comme le montre le dépôt XLA sur GitHub. Parallèlement, le projet TorchTPU intègre des spécificités matérielles, incluant des noyaux ATen personnalisés pour les TPU de Google, permettant par exemple le support du parallélisme de données (Data Parallelism) et du parallélisme de tenseurs (Tensor Parallelism) pour des modèles comme Llama ou Qwen, disponibles dans les dépôts associés sur GitHub.

Le développement de ces infrastructures repose sur des outils comme Bazel pour la gestion des dépendances. L’écosystème s’ouvre au-delà des centres de données de Google. Oracle, par exemple, a annoncé en octobre 2026 qu’il s’associerait à Google pour proposer ses TPU via Oracle Cloud Infrastructure à partir de 2027, dans une stratégie visant à diversifier les fournisseurs face à la dépendance à Nvidia (The Stack).

Le projet TorchTPU est en phase de maturation et son feuille de route inclut des travaux sur les formes dynamiques, l’amélioration des performances de compilation, l’extension de la couverture des opérateurs et l’intégration du renforcement apprenant.

Sources

• GitHub – openxla/xla: A machine learning compiler for GPUs, CPUs, and ML accelerators, github.com
• GitHub – google-pytorch/torch_tpu, github.com
• Google and Meta Open-source TorchTPU to Bring Native PyTorch Support to Google TPUs, InfoQ, 2 novembre 2026
• Oracle taps Google for TPUs as hyperscalers hedge Nvidia exposure, The Stack, 28 octobre 2026

Crédits photo

• Image en une : Google Data Center, The Dalles, par Visitor7, CC BY-SA 3.0, via Wikimedia Commons (recadrée)
• Image dans l’article : Google Mayes County P0004991a, par Xpda, CC BY-SA 4.0, via Wikimedia Commons (recadrée)