PyTorch sur TPU et Trainium : jusqu’où réduire la dépendance Nvidia ?

Faire tourner PyTorch sur TPU et Trainium ne demande plus qu’un changement d’appareil avec TorchTPU et TorchNeuron. Chiffres, limites, coûts réels et qui doit tester.

Actualité

Faire tourner PyTorch sur les TPU de Google ou les puces Trainium d’AWS demandait jusqu’ici une bibliothèque intermédiaire et pas mal d’adaptations. Les 19 et 20 octobre 2026, juste avant la PyTorch Conference, les deux groupes ont ouvert un backend PyTorch natif : TorchTPU (Google, développé avec l’équipe PyTorch de Meta) et TorchNeuron (AWS, en bêta publique). Objectif commun : que le code PyTorch écrit pour les GPU Nvidia fonctionne sur leurs puces avec très peu de modifications.

• Changer d’appareil plutôt que réécrire. Pour TorchNeuron, AWS indique que le seul changement est model.to('neuron'). Pour TorchTPU, Google parle de « modifications minimes » ; dans Hugging Face Transformers v5, il suffit de device_map="tpu".
• Les deux modes d’exécution. Les deux gèrent le mode eager (opération par opération, facile à déboguer) et torch.compile (compilation du graphe complet pour la production).
• Entraînement distribué intégré. DDP, FSDP2 et DTensor fonctionnent sur les deux ; TorchNeuron ajoute le parallélisme de tenseurs et de pipeline.
• Branchés sur l’écosystème d’inférence. TorchTPU fournit des backends pour vLLM et SGLang ; TorchNeuron propose vLLM Neuron et un support dans Transformers v5, Diffusers et trl.
• Open source sous licence Apache 2.0 pour les deux projets sur GitHub.

Analyse

Analyse : une ligne de code suffit-elle pour quitter les GPU Nvidia ?

Réponse courte : assez pour tester, pas pour migrer. Le vrai changement tient au mode d’intégration ; les économies dépendent fortement des modèles que vous exécutez.

Qu’est-ce qui change par rapport à avant ?

La différence clé, c’est PrivateUse1, le point d’extension officiel de PyTorch pour le nouveau matériel. Un tenseur sur TPU ou Trainium est un tenseur PyTorch normal, sans enveloppe. C’est pourquoi des bibliothèques comme Hugging Face, vLLM ou TorchTitan n’ont plus besoin d’une branche par puce.

Pourquoi cela tourne-t-il sans réécriture ?

Quand vous appelez torch.compile, PyTorch utilise TorchDynamo pour capturer le graphe de calcul. Au lieu de le confier à un compilateur GPU, TorchTPU le traduit en StableHLO et laisse XLA l’optimiser pour TPU. TorchNeuron suit un chemin proche via Torch-MLIR, StableHLO et le compilateur Neuron. Le développeur continue d’écrire du PyTorch ; la traduction vers la puce se fait en dessous.

Le mode « Fused Eager » de TorchTPU regroupe automatiquement de nombreuses petites opérations en blocs plus grands avant de les envoyer à la puce. Google a annoncé des gains de 50 % à plus de 100 % par rapport à l’exécution opération par opération. C’est un chiffre du fabricant, sans mesure indépendante.

Qui est concerné, et comment ?

• Ingénieurs ML : peuvent exécuter la même boucle d’entraînement sur GPU, TPU et Trainium pour comparer. Les kernels maison exigent toujours d’apprendre Pallas (TPU) ou NKI (Trainium).
• Responsables techniques : gagnent un levier pour négocier les prix des GPU. Un contrat cloud n’est plus lié à CUDA si les modèles tournent sur plusieurs puces.
• PME : la plupart utilisent des modèles via API, l’effet est donc indirect. Les équipes qui hébergent des modèles ouverts (Qwen, Llama) sur Google Cloud ou AWS en profitent directement.
• Marketing et produit : rien ne change tout de suite. Le bénéfice éventuel viendra de prix d’API plus bas chez les fournisseurs utilisant ces puces.

Combien cela coûte-t-il réellement ?

Aucune des deux entreprises n’a publié de nouveaux tarifs. Le seul chiffre de coût est celui de Deepgram : environ 60 % d’audio traité en plus par dollar qu’avec ses GPU actuels. Le tableau est un exemple hypothétique appliquant ce chiffre à deux échelles, avec un coût de migration supposé de 30 000 USD (environ deux ingénieurs pendant 1 à 2 mois).

À retenir : avec une petite facture, changer de puce est rarement rentable. Avec une grosse facture et un modèle stable, le calcul commence à fonctionner. Les 60 % viennent d’un seul client et d’une seule charge (audio) ; votre modèle peut faire mieux, moins bien, ou ne rien gagner.

Risques et limites

• Couverture des opérateurs incomplète. Lors des tests PyTorch d’août 2026, environ 22 000 tests ont été ignorés car TorchNeuron n’avait pas l’opérateur correspondant. La première version ne gère ni les formes dynamiques, ni les tenseurs creux, ni NestedTensor ; le FP8 est en bêta privée.
• Les formes dynamiques restent un point faible sur TPU. Google reconnaît que des longueurs de séquence ou tailles de lot variables provoquent des recompilations et que la première itération a une latence élevée. Un chatbot aux prompts de longueur très variable est exactement la charge concernée.
• La dépendance ne fait que se déplacer. Les TPU n’existent que sur Google Cloud, Trainium que sur AWS. Quitter CUDA peut signifier se lier davantage à un cloud.
• Le code propre à CUDA ne tourne pas. CUDA graphs, appels spécifiques à NCCL et kernels CUDA maison doivent être réécrits.

Points à vérifier

• Les 52 % de MFU et les 8 192 puces proviennent de tests internes de Google avec TorchTitan, sans comparaison indépendante avec des GPU de prix équivalent.
• Les chiffres de Zenteiq (environ ×4) et de Deepgram (+60 % par dollar) sont des témoignages clients choisis par les fabricants.
• Le chiffre « 85 % des applications d’IA utilisent PyTorch » cité par TechCrunch n’a pas de source nommée ; TBN Grow n’a pas pu le vérifier.
• Pour vérifier vous-même : prenez un modèle en production, mesurez tokens/seconde et coût par million de tokens sur vos GPU, puis exécutez le même modèle sur TPU v6e ou Trainium2 avec les mêmes longueurs d’entrée.

L’avis de TBN Grow

C’est une évolution du logiciel d’infrastructure, utile pour les équipes qui entraînent ou hébergent des modèles ouverts sur Google Cloud ou AWS. Celles dont la facture GPU atteint plusieurs dizaines de milliers de dollars par mois devraient tester dès maintenant sur un modèle stable. Petites équipes, utilisateurs d’API ou charges aux entrées très variables gagneront à attendre une ou deux versions, le temps que la couverture des opérateurs et les formes dynamiques mûrissent.

À suivre

Ce qu’il faut surveiller pour PyTorch sur TPU et Trainium

Ce que vous pouvez faire dès maintenant

• [ ] Repérez chaque appel direct à cuda dans votre code et remplacez-le par l’API torch.accelerator pour faciliter le changement de puce.
• [ ] Testez TorchTPU : pip install torch_tpu==0.2.0 sur une VM TPU v6e (le README du projet la présente comme le démarrage le plus simple ; ajoutez --extra-index-url https://download.pytorch.org/whl/cpu pour éviter la version CUDA).
• [ ] Testez TorchNeuron : pip install torch-neuronx sur une instance trn2, puis passez l’appareil à neuron.
• [ ] Lancez les tests de votre propre projet pour repérer les opérateurs non pris en charge avant d’estimer les coûts.
• [ ] Comparez le coût par million de tokens ou par heure d’audio, pas seulement la vitesse.

Calendrier et ce qui n’est pas encore annoncé

• AWS prévoit de rendre TorchAO et le FP8 disponibles en version générale à la prochaine sortie et de combler les manques d’opérateurs ; sans date.
• Google vise la prise en charge des formes dynamiques dans torch.compile et une bibliothèque de kernels précompilés ; sans calendrier.
• Google n’a qualifié TorchTPU ni de bêta ni de version générale, et n’a pas publié la liste complète des générations de TPU prises en charge.
• Nvidia n’a pas réagi officiellement.

À lire aussi : Le marché chinois des puces IA bascule : Nvidia à 34 %, Huawei à 28 % et la question de la rentabilité des investissements dans l’infrastructure IA.

Sources de vérification : Google Cloud Blog, Google Developers Blog, TorchTPU sur GitHub, PyTorch Blog – TorchNeuron, AWS What's New, TechCrunch.