Entraîner un modèle sur un TPU à hôte unique à l'aide de Pax


Ce document présente brièvement comment travailler avec Pax sur un TPU à hôte unique (v2-8, v3-8, v4-8).

Pax est un framework permettant de configurer et exécuter des tests de machine learning en plus de JAX. Mise au point Pax sur la simplification du ML à grande échelle en partageant des composants d'infrastructure à l'aide de frameworks de ML existants et d'utiliser Bibliothèque de modélisation Praxis pour plus de 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 à utiliser votre projet Cloud TPU et d'installer les composants nécessaires à l'entraînement un modèle exécutant Pax sur un TPU à un seul hôte.

Installer Google Cloud CLI

La CLI Google Cloud contient des outils et des bibliothèques permettant d'interagir avec les produits et services de la CLI Google Cloud. Si vous ne l'avez pas déjà installé, faites-le maintenant en suivant les instructions de la section Installer la Google Cloud CLI.

Configurer la commande gcloud

(Exécutez gcloud auth list pour afficher vos 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 Google Cloud Console.

$ 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 vous connectez directement à la VM TPU via SSH. Vous pouvez exécuter des fonctions coder, 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 d'un environnement Cloud Shell ou du terminal d'ordinateur sur lequel la Google Cloud CLI est installée.

Définissez les zone en fonction de la disponibilité dans votre contrat. référence 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 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 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

Assurez-vous que tout est correctement installé 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 s'affiche. Il doit être égal à 8 si vous utilisez un v2-8, v3-8 ou 4 si vous utilisez une version v4-8.

Exécuter du code Pax sur une VM TPU

Vous pouvez maintenant exécuter n'importe quel code Pax. lm_cloud exemples sont un excellent point de départ pour exécuter des modèles dans Pax. Par exemple, les commandes suivantes permettent d'entraîner Paramètre 2B : modèle de langage SPMD basé sur Transformer sur des données synthétiques.

Les commandes suivantes affichent le résultat de l'entraînement pour un langage SPMD du modèle. L'entraînement dure 300 étapes en environ 20 minutes.

(vm)$ python3 .local/lib/python3.10/site-packages/paxml/main.py  --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir

Sur la tranche v4-8, la sortie doit inclure:

Pertes et temps de pas

Tensor de résumé à l'étape =step_# loss =loss
Tensor de résumé à l'étape =step_# Pas par seconde 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 ce n'est 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

Étape suivante

Pour en savoir plus sur Cloud TPU, consultez les pages suivantes :