Google rend TorchTPU open source, un backend natif PyTorch pour les Cloud TPU

Google a rendu open source TorchTPU, un backend natif PyTorch pour les Cloud TPUs. C'était co-développé avec l'équipe PyTorch de Meta et annoncé à la PyTorch Conference 2026 le 20 octobre 2026. Le code PyTorch existant cible un TPU via un torch.device("tpu"). Les TPU s'ajoutent désormais aux GPU et aux CPU en tant que périphérique PyTorch natif. L'équipe Google indique que TorchTPU remplace PyTorch/XLA et le recommande pour toutes les nouvelles charges de travail.

TL;DR

• Taille : Une pile logicielle, pas un modèle. 4 dépôts ouverts : TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan. TorchTPU est en v0.2.0 sur PyPI, Python 3.11+, Apache 2.0.
• Fonctionne sur : les TPU de Google Cloud. Testé sur des générations à partir de la v5e, jusqu'à 8,192 puces Ironwood (TPU7x). Le serving prend en charge les checkpoints FP8, MXFP4 et NVFP4.
• Performances : Mode eager pour le débogage, torch.compile pour un débit de niveau production.
• Le meilleur : 52% de MFU pour l'entraînement de Qwen 3 8B avec TorchTitan sur Ironwood (auto-déclaré).
• Le pire : Le mode eager prêt à l'emploi prenait 11.49s par étape dans un cas de fine-tuning, contre 2.78s entièrement optimisé.

Qu'est-ce que TorchTPU ?

TorchTPU est un backend PyTorch open source de Google et Meta qui exécute nativement les modèles PyTorch sur les TPU. Il s'intègre via l' interface PrivateUse1de PyTorch. Les tenseurs sont des tenseurs PyTorch ordinaires, sans sous-classes ni wrappers. Il exécute un processus par périphérique et préserve les vues de tenseurs et les mutations en place.

Son prédécesseur, PyTorch/XLA, enregistrait les graphes de lazy tensors. Les développeurs ajoutaient souvent du code pour marquer les régions compilables. PyTorch/XLA a aussi ne prenait en charge que le code purement SPMD. TorchTPU gère également l'exécution MPMD divergente.

Comment TorchTPU exécute-t-il le code PyTorch ?

TorchTPU est disponible 3 modes eagerly:

• Strict (par défaut) : distribue chaque opérateur de manière asynchrone avec la sémantique standard de PyTorch.
• Débogage : s'exécute de manière synchrone avec la vérification des bornes des tenseurs.
• DeferAndFuse : fusionne les séquences d'opérateurs eager pour réduire le trafic HBM. Le papier TorchTPU indique que cela peut être 50% ou plus plus rapide que Strict.

Pour une vitesse maximale, torch.compile(backend="tpu") capture le graphe avec TorchDynamo. Il abaisse via StableHLO vers OpenXLA, qui émet du code machine TPU. Les exécutables compilés sont mis en cache entre les processus, les hôtes et les redémarrages de tâches.

Les noyaux Pallas personnalisés sont appelables comme des fonctions PyTorch standard. Les noyaux Tokamax préconstruits et les noyaux de tuiles Helion fonctionnent également.

La distribution utilise le tpu_dist groupe de processus. Il prend en charge DDP, FSDP2 et DTensor, ainsi qu'un mesh sensible à la topologie pour les interconnexions TPU.

Quelle est la vitesse de TorchTPU ?

Tous les chiffres ci-dessous sont auto-déclarés par Google et ses partenaires.

• Échelle d'entraînement : TorchTitan a atteint jusqu'à 52% de MFU sur Qwen 3 8B. Google rapporte un scaling quasi linéaire jusqu'à 8,192 TPUs Ironwood.
• Échelle de réglage : Dans un cas de fine-tuning TorchTitan, le temps de pas est passé de 11.49s à 2.78s. DeferAndFuse à lui seul a donné 1.15x. Full torch.compile a donné 3.80x.
• Couverture de modèles : Google teste en continu environ 5,000 modèles Hugging Face en mode eager et compilé.
• Service de flotte : Zenteiq a fait passer la production quotidienne de 1.47B à 5.97B tokens pour 1,000 puces TPU v6e. Ce gain est venu du routage llm-d sur GKE, et non de TorchTPU seul.

Quelles intégrations sont livrées avec TorchTPU ?

• Hugging Face : Transformers v5 se déploie sur TPU avec device_map="tpu". Diffusers prend en charge l'exécution eager, torch.compile et le parallélisme tensoriel.
• vLLM-TorchTPU: construit avec Inferact. Conserve les ordonnanceurs amont et les serveurs d'API, avec des noyaux Pallas personnalisés.
• SGLang-TorchTPU: construit avec RadixArk. Conserve la réutilisation de préfixe RadixAttention.
• Fonctionnalités de service : caches KV FP8, checkpoints FP8, MXFP4 et NVFP4. Le décodage spéculatif inclut MTP et EAGLE-3.
• TorchTitan: recettes pour Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE et FLUX.1 12B.
• Également : llm-d 1.0, Ray, TorchRec, VeRL (aperçu) et LiteRT-Torch pour l'export sur appareil.

Comment TorchTPU se compare-t-il à PyTorch/XLA, TorchNeuron et JAX ?

Comment installer TorchTPU ?

Installez le paquet sur une VM TPU, puis sélectionnez le tpu device :

# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)

Le README du dépôt épingle torch_tpu==0.2.0 avec l’index des wheels PyTorch CPU. Sans cela, pip récupère la build CUDA de torch. Les guides complets se trouvent sur torchtpu.readthedocs.io.

Points clés à retenir

• TorchTPU est open source sous Apache 2.0, co-construit avec Meta.
• Il remplace PyTorch/XLA par une exécution eager-first basée sur PrivateUse1.
• torch.compile(backend="tpu") abaisse les coûts via StableHLO et XLA.
• TorchTitan a atteint 52% de MFU sur Qwen 3 8B, avec une montée en charge jusqu’à 8 192 puces Ironwood.
• vLLM, SGLang et Hugging Face fonctionnent désormais sur TPU via PyTorch standard.

Consultez le dépôt GitHub, article, documentation et l’ annonce de Google Cloud. Tout le mérite revient au chercheur de ce projet. N’hésitez pas également à nous suivre sur Twitter et n’oubliez pas de rejoindre notre SubReddit ML de plus de 150k et de vous abonner à notre newsletter. Attendez ! Êtes-vous sur telegram ? vous pouvez désormais nous rejoindre sur telegram également.

L’article Google rend TorchTPU open source, un backend natif PyTorch pour les Cloud TPU est apparu en premier sur MarkTechPost.