Exécuter des applications Cloud TPU dans un conteneur Docker

Les conteneurs Docker facilitent la configuration des applications en combinant votre code et toutes les dépendances nécessaires dans un seul package distribuable. Vous pouvez exécuter des conteneurs Docker dans des VM TPU pour simplifier la configuration et le partage de vos applications Cloud TPU. Ce document explique comment configurer un conteneur Docker pour chaque framework de ML compatible avec Cloud TPU.

Important :Vous pouvez exécuter au maximum 256 pods Docker sur une VM TPU.

Entraîner un modèle TensorFlow dans un conteneur Docker

Appareil TPU

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

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. Créer un bucket Cloud Storage

   gcloud storage buckets create gs://your-bucket-name --location=europe-west4

3. Créer une VM TPU

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

4. Copiez le Dockerfile sur votre VM TPU

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

5. Connection SSH à la VM TPU

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

6. Compiler l'image Docker

   sudo docker build -t your-image-name .

7. Démarrer le conteneur Docker

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

8. Définir des variables d'environnement

   export STORAGE_BUCKET=gs://your-bucket-name
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x

9. Entraîner ResNet

   python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"

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

Pod TPU

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

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. Créer une VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.18.0-pod-pjrt

3. Copiez le Dockerfile sur votre VM TPU

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

4. Connection SSH à la VM TPU

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

5. Compiler l'image Docker

   sudo docker build -t your-image-name .

6. Démarrer un conteneur Docker

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

7. Entraîner ResNet

   python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

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

Entraîner un modèle PyTorch 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 en SSH à la VM TPU

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

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

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash

   Remarque :Une fois cette commande exécutée, votre invite de commande change, indiquant que le terminal est connecté au conteneur en cours d'exécution.

4. 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 section "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"

5. Cloner le dépôt PyTorch XLA

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

6. 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

Pod TPU

Lorsque vous exécutez du code PyTorch sur un pod 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. Exécutez le script d'entraînement dans un conteneur sur tous les nœuds de calcul TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/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éfinir 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

Entraîner un modèle JAX dans un conteneur Docker

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 en SSH à la VM TPU

   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 -ti --rm --name your-container-name --privileged --network=host python:3.8 bash

5. Installer JAX

   pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

6. Installer FLAX

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

7. 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

Pod TPU

Lorsque vous exécutez du code JAX sur un pod 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.8
   RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   WORKDIR ./flax/examples/mnist

2. Compiler l'image Docker

   docker build -t your-image-name .

3. 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-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

4. Transférer votre image Docker dans Artifact Registry

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

5. 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

6. Extrayez l'image Docker à partir 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"

7. Exécutez 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:your-tag bash"

8. Exécutez le script d'entraînement sur tous les nœuds de calcul 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êtez le conteneur sur tous les nœuds de calcul:

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

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

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

Entraîner un modèle JAX dans un conteneur Docker à l'aide de la pile stable JAX

Important :La pile stable JAX est en phase expérimentale.

Vous pouvez créer les images Docker MaxText et MaxDiffusion à l'aide de l'image de base de la pile stable JAX.

La pile stable JAX fournit un environnement cohérent pour MaxText et MaxDiffusion en regroupant JAX avec des packages de base tels que orbax, flax, optax et libtpu.so. Ces bibliothèques sont testées pour garantir la compatibilité et fournir une base stable pour créer et exécuter MaxText et MaxDiffusion. Cela élimine les conflits potentiels dus à des versions de paquets incompatibles.

La pile stable JAX inclut une libtpu.so entièrement publiée et qualifiée, la bibliothèque principale qui gère la compilation, l'exécution et la configuration réseau ICI des programmes TPU. La version libtpu remplace le build quotidien précédemment utilisé par JAX et garantit la fonctionnalité cohérente des calculs XLA sur TPU avec des tests de qualification au niveau PJRT dans les IR HLO/StableHLO.

Pour créer l'image Docker MaxText et MaxDiffusion avec la pile stable JAX, lorsque vous exécutez le script docker_build_dependency_image.sh, définissez la variable MODE sur stable_stack et la variable BASEIMAGE sur l'image de base que vous souhaitez utiliser.

L'exemple suivant spécifie us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1 comme image de base:

   bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1
   

Pour obtenir la liste des images de base de la pile stable JAX disponibles, consultez la section Images de la pile stable JAX dans Artifact Registry.

Étape suivante

• Tutoriels Cloud TPU
• Gérer les TPU
• Architecture du système Cloud TPU
• Exécuter du code TensorFlow sur des tranches de pods Cloud TPU
• Exécuter du code JAX sur des tranches de pods TPU