Présentation du projet

Google a publié TorchTPU, un backend natif qui permet d’exécuter les modèles PyTorch directement sur les Tensor Processing Units (TPU). Le dépôt GitHub fournit des ATen kernels spécifiques aux TPU, une compilation cache pour optimiser les calculs XLA, ainsi que des exemples couvrant des modèles variés tels que Llama, Qwen, ResNet, DLRM, minGPT et UNet. Le projet s’appuie sur la compatibilité de version Python de PyTorch, supportant actuellement Python 3.11+.

Architecture et compilation

Le cœur de TorchTPU repose sur des kernels ATen écrits en C++ qui traduisent les opérations PyTorch en appels XLA/MLIR. La compilation cache stocke les graphes XLA générés, réduisant le temps de re‑compilation lors d’exécutions répétées. La construction des roues utilise bazel avec le flag --config=wheel_common, qui active le backend partagé libxla_base.so et le cache distant RBE. Sans ce flag, chaque version de PyTorch créerait une copie du backend, entraînant des conflits d’enregistrement statique au moment de l’import.

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

Les builds non optimisés (-c fastbuild ou -c dbg) sont explicitement découragés car ils dégradent les performances. Si un développeur a besoin d’un wheel de débogage, il doit ajouter --//:allow_unoptimized_wheel=True après le flag d’optimisation.

Support et utilisation

Le backend s’intègre à torch.compile(), offrant un mode graphe qui profite de la compilation XLA. Les utilisateurs peuvent lancer des entraînements distribués grâce à la prise en charge du Data Parallelism et du Tensor Parallelism. Les notebooks d’exemple, fournis au format Marimo, s’exécutent sur une VM TPU ou une instance GCE v6e, et sont accessibles via la commande marimo edit file.py. L’installation se fait en deux temps : création d’un environnement virtuel puis installation du package torch_tpu==0.2.0 depuis le dépôt PyPI, en précisant l’index CPU de PyTorch pour éviter la version CUDA.

python -m venv .venv
source .venv/bin/activate
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

Pour les développeurs qui souhaitent compiler leur propre wheel, le processus inclut la génération du compile_commands.json via le script tools/setup_clangd.py, puis l’utilisation de VSCode avec l’extension clangd pour la navigation C++.

Limites et perspectives

Bien que TorchTPU propose un chemin complet du code source à l’exécution, plusieurs contraintes subsistent. Le build dépend de Bazelisk et de l’accès à un cache RBE ; sans credentials RBE, il faut ajouter --config=no_rbe, ce qui peut ralentir la compilation initiale. De plus, la documentation indique que les dépendances Python « TBD », laissant une incertitude sur les versions exactes requises au-delà de la matrice de compatibilité PyTorch. Enfin, le support officiel se limite à Python 3.11+, excluant les environnements legacy. Ces points suggèrent que, malgré une intégration fonctionnelle, les équipes devront planifier du temps supplémentaire pour la configuration d’infrastructure et la validation des versions.