Exécuter des charges de travail TPU dans un conteneur Docker

Appareil TPU

1. Créer une VM Cloud TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. Se connecter à la VM TPU à l'aide de SSH

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

3. Assurez-vous que le rôle Lecteur du registre des artefacts a été attribué à votre utilisateur Google Cloud . Pour en savoir plus, consultez Attribuer des rôles Artifact Registry.

4. Démarrer un conteneur dans la VM TPU à l'aide de l'image PyTorch/XLA quotidienne

   sudo docker run --net=host -ti --rm --name your-container-name --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
   bash

5. Configurer l'environnement d'exécution TPU

   Il existe deux options d'exécution PyTorch/XLA: PJRT et XRT. Nous vous recommandons d'utiliser PJRT, sauf si vous avez une raison d'utiliser XRT. Pour en savoir plus sur les différentes configurations d'exécution, consultez la documentation d'exécution PJRT.

   PJRT

   export PJRT_DEVICE=TPU

   XRT

   export XRT_TPU_CONFIG="localservice;0;localhost:51011"

6. Cloner le dépôt PyTorch XLA

   git clone --recursive https://github.com/pytorch/xla.git

7. Entraîner ResNet50

   python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

Une fois le script d'entraînement terminé, veillez à nettoyer les ressources.

1. Saisissez exit pour quitter le conteneur Docker.
2. Saisissez exit pour quitter la VM TPU.

3. Supprimer la VM TPU

   gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Tranche TPU

Lorsque vous exécutez du code PyTorch sur une tranche TPU, vous devez exécuter votre code sur tous les nœuds de calcul TPU en même temps. Pour ce faire, vous pouvez utiliser la commande gcloud compute tpus tpu-vm ssh avec les options --worker=all et --command. La procédure suivante vous explique comment créer une image Docker pour faciliter la configuration de chaque nœud de travail TPU.

1. Créer une VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-32 \
   --version=tpu-ubuntu2204-base

2. Ajouter l'utilisateur actuel au groupe Docker

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command='sudo usermod -a -G docker $USER'

3. Cloner le dépôt PyTorch XLA

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="git clone --recursive https://github.com/pytorch/xla.git"

4. Exécuter le script d'entraînement dans un conteneur sur tous les nœuds TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -v ~/xla:/xla -e PJRT_DEVICE=TPU us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 python /xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

   Options de commande Docker:

   • --rm supprime le conteneur une fois son processus terminé.
   • --privileged expose l'appareil TPU au conteneur.
   • --net=host lie tous les ports du conteneur à la VM TPU pour permettre la communication entre les hôtes du pod.
   • -e définit des variables d'environnement.

Une fois le script d'entraînement terminé, veillez à nettoyer les ressources.

Supprimez la VM TPU à l'aide de la commande suivante:

gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=us-central2-b

Appareil TPU

1. Créez la VM TPU.

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. Se connecter à la VM TPU à l'aide de SSH

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a

3. Démarrer le daemon Docker dans la VM TPU

   sudo systemctl start docker

4. Démarrer le conteneur Docker

   sudo docker run --net=host -ti --rm --name your-container-name \
   --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
   bash

5. Installer JAX

   pip install jax[tpu]

6. Installer FLAX

   pip install --upgrade clu
   git clone https://github.com/google/flax.git
   pip install --user -e flax

7. Installer les packages tensorflow et tensorflow-dataset

   pip install tensorflow
   pip install tensorflow-datasets

8. Exécutez le script d'entraînement FLAX MNIST.

   cd flax/examples/mnist
   python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5

Une fois le script d'entraînement terminé, veillez à nettoyer les ressources.

1. Saisissez exit pour quitter le conteneur Docker.
2. Saisissez exit pour quitter la VM TPU.

3. Supprimer la VM TPU

   gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Tranche TPU

Lorsque vous exécutez du code JAX sur une tranche TPU, vous devez exécuter votre code JAX sur tous les nœuds de travail TPU en même temps. Pour ce faire, vous pouvez utiliser la commande gcloud compute tpus tpu-vm ssh avec les options --worker=all et --command. La procédure suivante vous explique comment créer une image Docker pour faciliter la configuration de chaque nœud de travail TPU.

1. Créez un fichier nommé Dockerfile dans votre répertoire actuel et collez-y le texte suivant :

   FROM python:3.10
   RUN pip install jax[tpu]
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   RUN pip install tensorflow
   RUN pip install tensorflow-datasets
   WORKDIR ./flax/examples/mnist

2. Préparer un dépôt Artifact Registry

   gcloud artifacts repositories create your-repo \
   --repository-format=docker \
   --location=europe-west4 --description="Docker repository" \
   --project=your-project
   
   gcloud artifacts repositories list \
   --project=your-project
   
   gcloud auth configure-docker europe-west4-docker.pkg.dev

3. Compiler l'image Docker

   docker build -t your-image-name .

4. Ajoutez un tag à votre image Docker avant de la transférer vers Artifact Registry. Pour en savoir plus sur l'utilisation d'Artifact Registry, consultez Utiliser des images de conteneurs.

   docker tag your-image-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

5. Transférer votre image Docker dans Artifact Registry

   docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

6. Créer une VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

7. Extraire l'image Docker d'Artifact Registry sur tous les nœuds de travail TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command='sudo usermod -a -G docker ${USER}'

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"

8. Exécuter le conteneur sur tous les nœuds de travail TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag bash"

9. Exécuter le script d'entraînement sur tous les nœuds TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"

Une fois le script d'entraînement terminé, veillez à nettoyer les ressources.

1. Arrêter le conteneur sur tous les nœuds de travail

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"

2. Supprimer la VM TPU

   gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=europe-west4-a