AIREITERPasser de cuda à tpu dans un projet PyTorch ne demande souvent que quelques modifications avec TorchTPU. En pratique, l’accès au TPU, le choix des paquets, la compilation et la forme des tenseurs déterminent surtout si le portage fonctionnera correctement et offrira de bonnes performances.
Ce que TorchTPU apporte réellement
TorchTPU est un backend PyTorch natif pour les TPU Google. Ce n’est ni une API de modèles hébergée ni un remplacement d’une TPU VM. L’annonce de PyTorch met en avant l’exécution eager, torch.compile, les primitives distribuées et la prise en charge d’environ 450 opérations ATen. Le dépôt public recense également les kernels ATen, la mise en cache de la compilation, des exemples distribués et l’intégration de torch.compile.
Le principe de migration reste volontairement familier :
# Avant
model.to("cuda")
# Premier portage vers TPU
model.to("tpu")Mais cette modification vérifie surtout la portabilité du code, pas ses performances. PyTorch précise que le rendement maximal d’un TPU dépend encore des dimensions du modèle, du sharding et des opérations adaptées à ce matériel.
À préparer avant l’installation
TorchTPU n’est utile que si un TPU Google est effectivement disponible. Le dépôt recommande d’exécuter ses notebooks sur une TPU VM et présente une instance Google Compute Engine v6e comme le point de départ le plus simple. Un ordinateur portable dépourvu de TPU peut installer les paquets, mais il ne permettra pas de vérifier l’exécution sur le matériel.
Le README public prend en charge Python 3.11+ et indique la version torch_tpu==0.2.0. Utilisez un environnement isolé afin que l’expérience n’hérite pas d’une build CUDA ou de dépendances sans rapport.
Installer torch_tpu sans tirer CUDA
Créez un environnement virtuel, puis installez la roue CPU de PyTorch avec la version fixe de TorchTPU :
python -m venv .venv
source .venv/bin/activate
pip install torch_tpu==0.2.0 \
--extra-index-url https://download.pytorch.org/whl/cpuComme l’indique le README, l’absence de l’index supplémentaire peut entraîner l’installation d’une roue PyTorch CUDA, qui n’est pas la bonne dépendance pour une TPU VM.
Pour un premier portage, privilégiez la roue publiée. Les builds depuis les sources nécessitent Bazel ou Bazelisk, ainsi que les options documentées --config=wheel_common et, si nécessaire, --config=no_rbe.
Vérifier le backend avec un programme PyTorch minimal
Exécutez ce code sur la TPU VM :
import torch
x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y
print(z)
print(z.device)Si le résultat s’affiche avec un périphérique tpu, l’allocation des tenseurs et la multiplication matricielle sont bien passées par le backend TPU. Si le périphérique est indisponible, vérifiez d’abord le runtime de la VM et l’installation des paquets, avant de modifier le code du modèle.
TorchTPU prend en charge l’exécution eager pour le débogage et torch.compile pour l’exécution compilée. D’après l’annonce de PyTorch, son mode DeferAndFuse s’est montré au moins 50 % plus rapide que le mode eager strict dans les tests du projet. Il s’agit toutefois d’un résultat annoncé par le projet, pas d’une référence valable pour tous les modèles.
Porter un vrai workload PyTorch
Commencez par la modification minimale qui fonctionne, puis ajoutez la configuration distribuée propre aux TPU une fois le parcours mono-périphérique validé.
- Choisir le périphérique TPU : remplacez la sélection du périphérique par
torch.device("tpu")ou utiliseztorch.accelerator.current_accelerator()lorsque cette fonction est prise en charge. - Uniformiser les périphériques : déplacez le modèle et les tenseurs d’entrée vers le même périphérique.
- Commencer par un lot court : exécutez un seul batch avant d’activer la compilation ou l’exécution distribuée.
- Stabiliser les dimensions : gardez des tailles de batch et des formes de tenseurs constantes pendant les mesures ; un changement de forme peut relancer la compilation.
- Tester le parcours représentatif : validez le chemin réel du modèle et le mode distribué utilisé avant de lancer une migration plus large.
Un extrait CUDA distribué s’appuie généralement sur NCCL. L’exemple publié par TorchTPU utilise à la place le backend de groupe de processus TPU :
import torch
import torch.distributed as dist
device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)La commande exacte de lancement distribué dépend de l’environnement TPU et du workload. Modifier la chaîne du périphérique ne configure pas automatiquement le sharding, le placement des processus ni les pipelines d’entrée.
La première exécution ne ressemblera pas forcément aux suivantes. Le cache de compilation documenté par TorchTPU peut réutiliser les graphes répétés : mesurez donc séparément le temps de chauffe et le débit en régime établi.
Les cas déjà exploitables — et les zones encore perfectibles
Google documente des intégrations de TorchTPU avec vLLM, SGLang, TorchTitan et les outils Hugging Face. Son annonce consacrée à l’inférence décrit des parcours TPU natifs pour vLLM et SGLang, notamment la gestion du déplacement du KV cache et les fonctions de serving. L’annonce Google Cloud rapporte un résultat de 52 % de MFU avec Qwen 3 8B dans TorchTitan, ainsi que des cas d’utilisation en production.
Dans la pratique, les limites sont moins lisses que ne le laissent entendre les titres des annonces :
| Domaine | Ce que les sources établissent | Ce que vous devez vérifier |
|---|---|---|
| Installation | Python 3.11+, torch_tpu==0.2.0, index PyTorch CPU | La résolution correcte des dépendances avec votre image de VM et votre roue |
| Portabilité | Mode eager, torch.compile, périphérique tpu et prise en charge distribuée | Les opérateurs non pris en charge et les replis propres à chaque modèle |
| Performances | Environ 450 opérations ATen et un résultat annoncé de 52 % de MFU pour Qwen 3 8B dans TorchTitan | Le temps de chauffe, les formes, les tailles de batch et le coût sur votre modèle |
| Workflow développeur | Exemples publics, notebooks, outils de débogage et utilitaires de benchmark | L’accès au TPU, la complexité de Bazel et les possibilités d’observabilité |
| Limites actuelles | La couverture fonctionnelle continue d’évoluer | La dépendance éventuelle de votre stack à des fonctions qui ne sont pas encore documentées pour votre workload |
Quand adopter TorchTPU ?
TorchTPU mérite d’être testé dès maintenant si vous disposez déjà d’un accès à des TPU Google, utilisez des modèles PyTorch et souhaitez conserver un framework commun pour l’entraînement ou l’inférence. Le cas le plus favorable concerne les équipes prêtes à profiler et à ajuster les dimensions, plutôt que celles qui attendent d’un simple changement de périphérique une efficacité équivalente à celle d’un GPU.
Mieux vaut patienter avant de migrer un workload critique si vous avez besoin d’une matrice de compatibilité mature, d’un comportement prévisible au démarrage à froid ou d’une fonction indispensable que TorchTPU n’expose pas encore. Procédez par étapes : un modèle, une tranche TPU, un jeu d’évaluation fixe, puis des mesures distinctes pour le temps d’installation, la première compilation, le débit en régime établi et la correction des résultats.
FAQ TorchTPU
TorchTPU est-il identique à PyTorch/XLA ?
Non. TorchTPU est un backend PyTorch natif plus récent, conçu autour de l’exécution eager et de torch.compile. PyTorch/XLA est l’ancienne voie d’intégration des TPU. Le code existant basé sur XLA peut donc nécessiter une analyse de migration distincte.
Puis-je installer TorchTPU sur un ordinateur portable sans TPU ?
Vous pouvez installer les paquets, mais pas valider l’exécution sur TPU sans une infrastructure Google TPU compatible. La configuration pratique du dépôt officiel part du principe que vous utilisez une TPU VM.
Pourquoi ma première étape TorchTPU est-elle lente ?
La première exécution peut compiler un graphe et remplir les caches. Mesurez séparément les itérations suivantes et conservez des formes stables pendant l’analyse des recompilations.
Remplacer cuda par tpu garantit-il les mêmes performances ?
Non. Le remplacement du périphérique valide la correction du portage, mais il faut profiler les formes, le sharding, la couverture des opérateurs et le temps de chauffe avant de comparer les performances.
TorchTPU est-il prêt pour la production ?
Google fait état d’utilisations en production pour l’inférence et l’entraînement, mais les composants de l’écosystème et les fonctions disponibles n’ont pas tous le même niveau de maturité. Testez le modèle, le moteur de serving et le schéma distribué exacts avant de vous engager.
Installez la roue avec version fixe sur une TPU VM, exécutez la vérification en 2 × 2, puis comparez un workload déjà chaud avec votre parcours GPU actuel. Vous séparerez ainsi la portabilité du code, le temps de chauffe du compilateur, la couverture des opérateurs et l’économie réelle du matériel.