Google heeft TorchTPU als open source beschikbaar gesteld. Het project maakt het mogelijk om modellen en bewerkingen uit het machinelearningframework PyTorch uit te voeren op Google's Tensor Processing Units, de gespecialiseerde chips voor kunstmatige intelligentie.
TorchTPU bevat eigen ATen-kernels voor TPU's. ATen is de onderliggende bibliotheek waarop veel PyTorch-bewerkingen zijn gebaseerd. Ook biedt het project een compilatiecache, waarmee eerder gecompileerde XLA-berekeningen opnieuw kunnen worden gebruikt om de uitvoering te versnellen.
Training en modeluitvoering
De software ondersteunt zowel gedistribueerde training als modelinferentie. Daarbij zijn technieken beschikbaar voor data parallelism en tensor parallelism. Die maken het mogelijk om berekeningen over meerdere apparaten te verdelen.
Google levert voorbeelden voor verschillende modellen en toepassingen. Daaronder vallen Llama, Qwen, ResNet, DLRM, minGPT en UNet. Ook werkt TorchTPU samen met torch.compile(), de PyTorch-functie die modellen kan omzetten naar een geoptimaliseerde uitvoeringsgrafiek.
Voor ontwikkelaars zijn daarnaast hulpmiddelen voor foutopsporing en benchmarking opgenomen. Een gebruikershandleiding en interactieve notebooks geven voorbeelden van het installeren en gebruiken van de software. De notebooks zijn gemaakt met Marimo, een omgeving die vergelijkbaar is met Jupyter.
Installatie en licentie
De software kan worden geïnstalleerd als Python-pakket. De meest recente openbare versie is 0.2.0. TorchTPU ondersteunt momenteel Python 3.11 en nieuwere versies. Voor het bouwen vanuit de broncode gebruikt het project Bazel, een systeem voor het beheren en uitvoeren van softwarebuilds.
TorchTPU is vooral bedoeld voor gebruik op TPU-omgevingen. In de documentatie wordt een TPU-virtuele machine aanbevolen, waarbij een v6e-instantie van Google Cloud als eenvoudige optie wordt genoemd. Gebruikers kunnen het pakket rechtstreeks installeren of zelf een wheel-bestand bouwen.
Het project valt onder de Apache 2.0-licentie. Daardoor kunnen ontwikkelaars de code gebruiken, aanpassen en verder verspreiden binnen de voorwaarden van die licentie. Google nodigt ook externe ontwikkelaars uit om bij te dragen aan het PyTorch-op-TPU-project.
