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

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