Ce document présente brièvement comment utiliser Pax sur un TPU à hôte unique (v2-8, v3-8, v4-8).
Pax est un framework qui permet de configurer et d'exécuter des tests de machine learning en plus de JAX. Pax s'efforce de simplifier le ML à grande échelle en partageant des composants d'infrastructure avec des frameworks de ML existants et en utilisant la bibliothèque de modélisation Praxis pour sa modularité.
Objectifs
- Configurer des ressources TPU pour l'entraînement
- Installer Pax sur un TPU à hôte unique
- Entraîner un modèle SPMD basé sur Transformer à l'aide de Pax
Avant de commencer
Exécutez les commandes suivantes pour configurer gcloud afin qu'il utilise votre projet Cloud TPU et installer les composants nécessaires à l'entraînement d'un modèle exécutant Pax sur un TPU à hôte unique.
Installer Google Cloud CLI
Elle contient des outils et des bibliothèques permettant d'interagir avec les produits et services Google Cloud CLI. Si ce n'est pas déjà fait, installez-la maintenant en suivant les instructions de la section Installer la Google Cloud CLI.
Configurer la commande gcloud
(Exécutez gcloud auth list pour voir les comptes disponibles).
$ gcloud config set account account
$ gcloud config set project project-id
Activer l'API Cloud TPU
Activez l'API Cloud TPU à l'aide de la commande gcloud suivante dans Cloud Shell.
(Vous pouvez également l'activer à partir de la console Google Cloud).
$ gcloud services enable tpu.googleapis.com
Exécutez la commande suivante pour créer une identité de service (un compte de service).
$ gcloud beta services identity create --service tpu.googleapis.com
Créer une VM TPU
Avec les VM Cloud TPU, votre modèle et votre code s'exécutent directement sur la VM TPU. Vous accédez directement en SSH à la VM TPU. Vous pouvez exécuter du code arbitraire, installer des packages, afficher les journaux et déboguer du code directement sur la VM TPU.
Créez votre VM TPU en exécutant la commande suivante à partir de Cloud Shell ou du terminal de votre ordinateur sur lequel la Google Cloud CLI est installée.
Définissez zone en fonction de la disponibilité indiquée dans votre contrat, et référencez Régions et zones TPU si nécessaire.
Définissez la variable accelerator-type sur v2-8, v3-8 ou v4-8.
Définissez la variable version sur tpu-vm-base pour les versions de TPU v2 et v3, ou sur tpu-vm-v4-base pour les TPU v4.
$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version
Se connecter à la VM Google Cloud TPU
Connectez-vous en SSH à votre VM TPU à l'aide de la commande suivante :
$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone
Lorsque vous êtes connecté à la VM, l'invite de l'interface système passe de username@projectname à username@vm-name:
Installer Pax sur la VM Google Cloud TPU
Installez Pax, JAX et libtpu sur votre VM TPU à l'aide des commandes suivantes:
(vm)$ python3 -m pip install -U pip \
python3 -m pip install paxml jax[tpu]
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Vérifier le système
Vérifiez que tout est installé correctement en vérifiant que JAX voit les cœurs de TPU:
(vm)$ python3 -c "import jax; print(jax.device_count())"
Le nombre de cœurs de TPU est affiché. Il doit être égal à 8 si vous utilisez la version v2-8 ou v3-8, ou 4 si vous utilisez la version v4-8.
Exécuter du code Pax sur une VM TPU
Vous pouvez maintenant exécuter le code Pax de votre choix. Les exemples lm_cloud constituent un excellent point de départ pour exécuter des modèles dans Pax. Par exemple, les commandes suivantes entraînent un modèle de langage SPMD basé sur un paramètre 2B à partir de données synthétiques.
Les commandes suivantes affichent les résultats d'entraînement pour un modèle de langage SPMD. Il effectue un entraînement pour 300 pas en 20 minutes environ.
(vm)$ python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
Sur les tranches v4-8, la sortie doit inclure:
Défaites et durées de pas
Tensor de résumé à l'étape=step_# loss = loss
Tensor de résumé à l'étape=step_# Pas/s x
Effectuer un nettoyage
Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.
Lorsque vous avez fini d'utiliser votre VM TPU, procédez comme suit pour nettoyer vos ressources.
Si vous ne l'avez pas déjà fait, déconnectez-vous de l'instance Compute Engine:
(vm)$ exit
Supprimez votre Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Étapes suivantes
Pour en savoir plus sur Cloud TPU, consultez les pages suivantes :