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

Attention : Cette page s'applique à l'API Cloud TPU, qui n'est plus en développement actif. L'API Cloud TPU ne recevra que des correctifs de bugs et des mises à jour de sécurité. Pour bénéficier des dernières fonctionnalités et de la compatibilité avec les dernières versions de TPU, nous vous recommandons d'utiliser Compute Engine ou Google Kubernetes Engine pour gérer les ressources TPU. Pour en savoir plus, consultez Ressources TPU dans Compute Engine et À propos des TPU dans GKE. L'API Cloud TPU n'est pas compatible avec les TPU7x ni les versions ultérieures de TPU.

Les conteneurs Docker facilitent la configuration des applications en combinant votre code et toutes les dépendances nécessaires dans un package unique 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 PyTorch dans un conteneur Docker

Avant d'exécuter les commandes suivantes, créez les variables d'environnement suivantes :

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Descriptions des variables d'environnement

• PROJECT_ID : ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
• TPU_NAME : nom du TPU.
• ZONE : Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones TPU.
• ACCELERATOR_TYPE : le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
• RUNTIME_VERSION : Version logicielle du Cloud TPU.

Appareil TPU

1. Créez une VM Cloud TPU.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Connectez-vous à la VM TPU à l'aide de SSH :

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

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

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

   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
   

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

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

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

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

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

6. Clonez le dépôt PyTorch XLA.

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

7. Entraînez ResNet50.

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

Une fois le script d'entraînement terminé, nettoyez les ressources.

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

3. Supprimez la VM TPU.

   gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE
   

Tranche TPU

Lorsque vous exécutez du code PyTorch sur une tranche de TPU, vous devez exécuter votre code sur tous les nœuds de calcul TPU en même temps. Pour cela, 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 calcul TPU.

1. Créez une VM TPU.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Ajoutez l'utilisateur actuel au groupe Docker.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --worker=all \
   --command='sudo usermod -a -G docker $USER'
   

3. Clonez le dépôt PyTorch XLA.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="git clone --recursive https://github.com/pytorch/xla.git"
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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 la commande Docker :

   • --rm supprime le conteneur une fois le 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 les variables d'environnement.

Une fois le script d'entraînement terminé, nettoyez les ressources.

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

gcloud compute tpus tpu-vm delete $TPU_NAME \
--zone=$ZONE

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

Avant d'exécuter les commandes suivantes, créez les variables d'environnement suivantes :

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Descriptions des variables d'environnement

• PROJECT_ID : ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
• TPU_NAME : nom du TPU.
• ZONE : Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones TPU.
• ACCELERATOR_TYPE : le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
• RUNTIME_VERSION : Version logicielle du Cloud TPU.

Appareil TPU

1. Créez la VM TPU.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Connectez-vous à la VM TPU à l'aide de SSH :

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

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

   sudo systemctl start docker
   

4. Démarrez 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. Installez JAX.

   pip install jax[tpu]
   

6. Installez FLAX.

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

7. Installez les packages tensorflow et tensorflow-datasets.

   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é, nettoyez les ressources.

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

3. Supprimez la VM TPU.

   gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE
   

Tranche TPU

Lorsque vous exécutez du code JAX sur une tranche de TPU, vous devez exécuter votre code JAX sur tous les nœuds de calcul TPU en même temps. Pour cela, 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 calcul TPU.

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

   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éparez un dépôt Artifact Registry.

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

3. Compilez 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/$PROJECT_ID/your-repo/your-image-name:your-tag
   

5. Transférez l'image Docker dans Artifact Registry.

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

6. Créez une VM TPU.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

7. Extrayez l'image Docker depuis Artifact Registry vers tous les nœuds de calcul TPU.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command='sudo usermod -a -G docker ${USER}'
   

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

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

8. Exécutez le conteneur sur tous les nœuds de calcul TPU.

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

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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é, nettoyez les ressources.

1. Arrêtez le conteneur sur tous les nœuds de calcul.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker kill your-container-name"
   

2. Supprimer la VM TPU

   gcloud compute tpus tpu-vm delete $TPU_NAME \
   --zone=$ZONE
   

Étapes suivantes

• Tutoriels Cloud TPU
• Gérer des TPU
• Architecture Cloud TPU
• Exécuter du code JAX sur des tranches TPU