Exécuter des applications Cloud TPU dans un conteneur Docker

Les conteneurs Docker créent la configuration des applications en associant votre code et tous les outils les dépendances dans un seul package distribuable. Vous pouvez exécuter des conteneurs Docker Des VM TPU pour simplifier la configuration et le partage de vos applications Cloud TPU. Ce décrit comment configurer un conteneur Docker pour chaque framework de ML compatibles avec Cloud 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 la commande 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.17.0-pjrt
    
  4. Copier le Dockerfile dans 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 la commande 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.17.0-pod-pjrt
    
  3. Copier le Dockerfile dans 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 à une 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 de nuit.

    sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash
    
  4. Configurer l'environnement d'exécution TPU

    Il existe deux options d'exécution PyTorch/XLA: PJRT et XRT. Nous vous recommandons utiliser PJRT sauf si vous avez une raison d'utiliser XRT. Pour en savoir plus sur la différentes configurations d'exécution. Vous avez une bonne raison d'utiliser XRT. À Pour en savoir plus sur les différentes configurations 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"
    
  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 le code sur tous les TPU les nœuds de calcul en même temps. Pour ce faire, vous pouvez utiliser la commande gcloud compute tpus tpu-vm ssh avec les commandes --worker=all et Indicateurs --command. La procédure suivante vous montre comment créer un pour faciliter la configuration de chaque nœud de calcul 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 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 a défini les 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 à une 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 TPU les nœuds de calcul en même temps. Pour ce faire, vous pouvez utiliser gcloud compute tpus tpu-vm ssh avec les options --worker=all et --command. Les éléments suivants : vous explique comment créer une image Docker pour configurer chaque TPU pour chaque nœud de calcul.

  1. Créez un fichier nommé Dockerfile dans votre répertoire actuel et collez la commande 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 Travailler avec 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 d'Artifact Registry sur tous les nœuds de calcul 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 calcul 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
    

Étape suivante