PyTorch est le framework derrière la plupart des modèles d'IA, utilisé par environ 85 % des applications d'IA selon certaines estimations. Il a été conçu avec les GPU Nvidia en tête, l'une des raisons pour lesquelles l'emprise de Nvidia sur le calcul d'IA est si difficile à desserrer. Cette semaine, les deux plus grands fournisseurs de cloud ont bougé en même temps. Le 19 octobre, Amazon a sorti TorchNeuron pour ses puces Trainium, et le 20 octobre, lors de la PyTorch Conference 2026, Google a publié en open source TorchTPU pour ses Tensor Processing Units, développé avec l'équipe PyTorch de Meta.
Ce qui change dans votre code
Les deux projets se branchent sur PyTorch via son interface officielle pour le nouveau matériel : les tenseurs restent des tenseurs PyTorch ordinaires. Pour la plupart des modèles, passer sur TPU revient à choisir l'appareil.
# pip install torch_tpu (on a Cloud TPU VM)
import torch
import torch_tpu # registers the "tpu" device
device = torch.device("tpu")
model = MyModel().to(device)
batch = batch.to(device)
loss = model(batch).loss
loss.backward()
# Optional: compile for extra speed
model = torch.compile(model, backend="tpu")import torch
device = torch.device("cuda")
model = MyModel().to(device)
batch = batch.to(device)
loss = model(batch).loss
loss.backward()
model = torch.compile(model)Comment fonctionne TorchTPU
- Mode Strict (par défaut)
- Exécute chaque opération comme du PyTorch normal, de façon asynchrone. Le plus simple à déboguer.
- Mode Debug
- Synchrone, avec contrôles de limites, pour traquer les erreurs.
- DeferAndFuse
- Regroupe les opérations pour réduire le trafic mémoire ; Google annonce 15 % de gain dans son test de fine-tuning, et 50 % ou plus dans certains cas.
- torch.compile(backend="tpu")
- Capture tout le graphe et le compile via OpenXLA en code machine TPU, mis en cache entre les redémarrages.
- Entraînement distribué
- DDP, FSDP2 et DTensor via un process group TPU, pour des charges identiques ou divergentes.
- Kernels personnalisés
- Kernels Pallas, Tokamax et Helion quand il faut une vitesse optimisée à la main.
Exemple TorchTitan annoncé par Google. Secondes par étape d'entraînement.
TorchTPU, PyTorch/XLA et TorchNeuron comparés
TorchTPU
TPU de Google- Statut
- open source, v0.2.0, Apache 2.0
- Puces
- TPU v5e et suivantes, jusqu'à Ironwood
- Se comporte comme du PyTorch normal, eager par défaut
- Prise en charge de vLLM, SGLang et Transformers (device_map="tpu")
La voie recommandée par Google pour tout nouveau travail sur TPU.
PyTorch/XLA
TPU de Google, ancien- Statut
- en cours de remplacement par TorchTPU
- Enregistre le graphe de façon paresseuse, il faut donc indiquer ce qu'il faut compiler
- Ne gère que des calculs identiques sur chaque puce
Prévoyez une migration si vous l'utilisez aujourd'hui.
TorchNeuron
Amazon Trainium- Statut
- bêta publique (Neuron 2.33)
- Puces
- Trainium 2 et Trainium 3
- Eager plus torch.compile via le compilateur Neuron
- FSDP, DTensor, DDP et parallélisme de tenseurs ; service avec vLLM Neuron
Amazon contribue aussi du code à PyTorch pour le rendre moins dépendant de Nvidia.
Faut-il migrer ?
- 1
Regardez ce que vous exécutez vraiment
Le fine-tuning d'un modèle Hugging Face ou son service avec vLLM est le cas le plus simple. Les kernels CUDA maison sont le plus difficile.
- 2
Chiffrez la même tâche sur chaque puce
Comparez le coût par étape d'entraînement ou par million de tokens servis, pas le prix à l'heure.
- 3
Commencez par une petite charge réelle
Portez un job de fine-tuning ou d'inférence et comparez vitesse et résultats.
- 4
Gardez un code portable
Utilisez du code indépendant de l'appareil et torch.compile, pour que changer de puce reste un réglage.
Testez vos connaissances
1/3Comment cibler une TPU avec du code PyTorch grâce à TorchTPU ?
TorchTPU se branche sur l'interface officielle de PyTorch pour le nouveau matériel.
2/3Que devient PyTorch/XLA ?
Google recommande TorchTPU pour toutes les nouvelles charges sur TPU.
3/3Quelles puces TorchNeuron cible-t-il ?
Il est en bêta publique dans Neuron 2.33.
Vous avez 0 bonne(s) réponse(s) sur 3.
Pour les ingénieurs Python et ML : les nouveautés de scikit-learn 1.9, Python 3.15 et Apache Spark 4.2.
Sources : MarkTechPost sur TorchTPU (20 octobre 2026) ; TechCrunch, « Google has joined Amazon in attacking Nvidia's stronghold with AI app developers » (20 octobre 2026) ; le README de TorchTPU. Les chiffres de performance sont ceux de Google ; consultez la documentation officielle avant d'installer.
