TorchTPU : ce que demande la migration de PyTorch vers TPU

Google a annoncé TorchTPU le 20 octobre 2026. Ce backend open source, développé avec l’équipe PyTorch de Meta, permet d’exécuter du code PyTorch sur les TPU de Google. Pour un premier lancement, il faut installer torch_tpu et sélectionner le périphérique TPU ; les travaux distribués demandent aussi une configuration propre à cette plateforme.

Exécuter un modèle PyTorch sur TPU

Un backend est la couche logicielle qui relie PyTorch au matériel. Avec TorchTPU, une exécution simple peut utiliser le mode eager, qui lance les opérations au fur et à mesure. Google indique aussi que torch.compile peut compiler le modèle pour le TPU.

La documentation Google donne cette commande d’installation :

pip install torch_tpu

Dans le code, il faut sélectionner torch.device("tpu"). L’exemple PyTorch montre aussi comment remplacer .to("cuda") par .to("tpu"). Le changement de périphérique peut suffire à lancer un modèle compatible, mais les formes des tenseurs, le partitionnement du travail et certaines opérations peuvent nécessiter des ajustements pour tirer pleinement parti du TPU.

La documentation Google sur PyTorch pour Cloud TPU détaille l’installation et le choix du périphérique.

Les tâches distribuées nécessitent une configuration TPU

Une exécution distribuée répartit le travail entre plusieurs appareils. TorchTPU utilise le backend tpu_dist de torch.distributed et prévoit un processus par appareil. Le lancement doit aussi recevoir des informations propres à la topologie TPU et au slice builder, qui configure la tranche de TPU utilisée.

En pratique, le premier essai peut se limiter au choix du périphérique. Une tâche distribuée exige aussi de préparer son lancement, et un modèle conçu autour d’opérations ou d’un partitionnement adaptés au GPU peut demander du travail supplémentaire.

Ce que mesurent les tests publiés par le projet PyTorch

Le projet PyTorch rapporte environ 450 opérateurs ATen pris en charge, hors délégués structurels ou composites. Ses tests couvrent plusieurs critères distincts :

• Dans un ensemble de modèles sélectionnés pour leur popularité, 4 821 sur 4 917 ont réussi un test d’exécution sans plantage, soit 98 % selon le projet.
• Dans un autre ensemble, composé de modèles représentatifs, 638 sur 657 ont atteint la parité numérique avec une référence GPU, soit 97 % selon le projet.
• Pour l’évaluation du dynamisme borné, 834 des 839 LLM causaux ont satisfait le test de correction, soit 99 % selon le projet.

Ces résultats portent sur les ensembles et critères testés par le projet. Ils ne constituent pas une comparaison standardisée des performances d’un même modèle sur GPU et TPU. Le projet indique également que le dynamisme borné reste en développement.

Pour illustrer le passage à grande échelle, le projet rapporte un entraînement de Qwen3 MoE 235B sur une configuration de 8 192 puces TPU v7. Avec une taille de lot de 262 144 et une longueur de séquence de 4 096, le test a atteint 2 104 jetons par seconde et par TensorCore, avec 27,1 % de MFU et un pas d’entraînement de 31,2 secondes.

Réutiliser des poids enregistrés sur TPU

Un state_dict PyTorch contient les paramètres du modèle, notamment ses poids et ses biais. Pour un modèle compatible, ces paramètres peuvent être chargés dans son architecture puis associés au périphérique TPU. La fonction torch.load permet notamment de remapper le périphérique lors du chargement. La documentation PyTorch sur l’enregistrement et le chargement des modèles décrit cette démarche.

Réutiliser les poids ne revient pas à transférer d’un seul coup tout un environnement d’entraînement. Les paramètres du modèle et les autres états enregistrés, comme ceux d’un optimiseur, sont des éléments distincts du checkpoint.