PyTorch s’ouvre à TPUs et Trainium, sans effacer Nvidia

À la PyTorch Conference North America 2026, Google et AWS ont présenté deux backends natifs pour PyTorch : TorchTPU pour les TPU de Google et TorchNeuron pour les puces Trainium d’AWS. L’idée est simple, mais stratégique : permettre aux développeurs de déplacer plus facilement leurs charges de travail hors de l’écosystème Nvidia, qui reste dominant. Dans les faits, la portabilité progresse, sans que la promesse de performance “clé en main” soit encore tenue.

Deux backends pour rendre les puces maison plus attrayantes

TorchTPU est désormais disponible sur GitHub et vise les TPU de Google, tandis que TorchNeuron est entré en bêta publique pour Trainium2 et Trainium3, chez AWS. Les deux outils veulent faire de PyTorch un point d’entrée plus naturel vers des accélérateurs conçus en interne par les deux groupes. Selon le texte de Google, TorchTPU permet d’exécuter des charges PyTorch existantes sans dépendance JAX dans le runtime, avec des changements de code qualifiés de “trivial code changes”.

AWS pousse la même logique avec TorchNeuron. Le passage vers Trainium peut, selon l’entreprise, se limiter dans bien des cas à un simple changement vers .to('neuron') pour l’exécution eager, puis à la sélection du backend Neuron pour les modèles compilés. Cette approche compte car elle enlève une partie de la friction qui décourage souvent les équipes de tester d’autres puces que celles de Nvidia. Mais elle ne supprime pas la nécessité d’optimiser pour la cible choisie.

Le vrai sujet reste la performance, pas l’annonce

Les deux acteurs savent que le sujet ne se résume pas à une compatibilité de façade. Pour les workloads distribués, PyTorch indique que l’approche par défaut de TorchTPU, dite “Skinny Collectives”, peut entraîner une pénalité de performance de l’ordre de 30 à 40 % par rapport à une approche optimisée, car elle empêche le compilateur XLA de faire se chevaucher calcul et communication. Autrement dit, la portabilité est là, la fluidité parfaite non.

PyTorch recommande, pour les charges de production qui visent le meilleur débit, des ajustements supplémentaires et des piles plus spécialisées, comme Pallas kernels, vLLM-TorchTPU, SGLang-TorchTPU ou TorchTitan-TorchTPU. Côté AWS, la société précise que la plupart des équipes pourront essayer Trainium sans réécriture complète, mais laisse entendre que certaines charges demanderont davantage de modifications. Un cas cité dans la documentation a nécessité trois changements pour adapter le modèle Laplace 3B : l’usage du PyTorch natif pour le kernel GLA attention, une refonte de la communication de Muon et la compilation du backward séquentiel.

PyTorch devient aussi une bataille d’infrastructure

Le mouvement dépasse le simple choix des puces. Selon Amazon, 85% des chercheurs et développeurs en IA écrivent leur code en PyTorch, ce qui en fait une couche logicielle décisive pour les fournisseurs de matériel. Chris Lattner, vice-président exécutif de l’IA chez Qualcomm, a résumé l’enjeu en disant : “We all know that PyTorch has done an amazing job at unifying the modeling layer, making it easy to build and train models for a wide variety of different applications, but we still have the bottom part of the stack that the hardware vendors have to make for themselves. This includes the kernels, the compilers, and low-level system software.”

Dans ce contexte, Google, AWS, Qualcomm et d’autres cherchent à combler ce “bas de pile” qui comprend les kernels, les compilateurs et les outils d’entraînement distribué. Google a aussi annoncé des intégrations avec Hugging Face, vLLM, SGLang, llm-d, TorchTitan, Ray et TorchRec, ainsi que l’arrivée de LiteRT dans l’écosystème PyTorch. Le message est limpide : si le modèle attire l’attention, l’infrastructure qui l’exécute finit toujours par reprendre la main.

La montée en charge oblige PyTorch à se réorganiser

PyTorch lui-même encaisse cette croissance. Alban Desmaison, co-mainteneur du projet, a indiqué au public que plus de 600 000 projets GitHub reposent sur PyTorch et qu’environ 150 pull requests arrivent chaque jour. Mark Collier, directeur exécutif de la PyTorch Foundation, a de son côté annoncé une hausse de 60 % des contributeurs actifs sur un an, de 30 % des organisations participantes et de 40 % des pull requests. Ces chiffres disent moins une victoire qu’une tension : la popularité du framework pèse sur les tests, l’infrastructure et les opérations de publication.

Edward Yang, également co-mainteneur, a replacé le débat à son niveau le plus concret : “At the end of the day, what you care about is how long it takes for you to get the final checkpoint that you’re actually going to serve.” Les équipes veulent des briques plus petites, plus simples à profiler et plus faciles à réparer quand quelque chose casse sur le chemin critique. L’obsession du benchmark pur cède du terrain à un autre besoin, plus prosaïque : sortir un modèle exploitable sans passer ses nuits à déboguer des couches opaques.

Les usages business rappellent pourquoi ce socle compte

Shopify a montré à la conférence comment PyTorch sert aussi à des tâches métier bien bornées. L’entreprise a décrit un flux d’apprentissage continu fondé sur un jeu de données de référence, un évaluateur calé sur les jugements humains, puis une boucle où les échecs en production alimentent le réentraînement. Pour son application de profil client, Shopify est partie d’environ 29 000 exemples et a enrichi son modèle à chaque itération. “It’s important to keep in mind that human agreement is the ceiling,” a expliqué Cody Mazza-Anthony, ingénieur ML chez Shopify. “We’re not looking for perfect here. We’re looking for as good as human.”

Ce type d’usage montre pourquoi l’ouverture de PyTorch sur d’autres puces compte vraiment. Si les entreprises peuvent déplacer plus facilement leurs charges entre Nvidia, Google et AWS, elles gagnent du pouvoir de négociation, un peu de marge sur les coûts et davantage de liberté technique. Reste que l’ouverture n’efface pas les compromis : on peut simplifier le portage, mais pas abolir les lois de la physique, dommage pour les brochures.