Ejecuta aplicaciones de Cloud TPU en un contenedor de Docker.

Dispositivo de TPU

1. Crea un archivo llamado Dockerfile en tu directorio actual y pega el texto siguiente

   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. Crea un bucket de Cloud Storage

   gsutil mb -c standard -l europe-west4 gs://your-bucket-name
   

3. Crea una VM de 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. Copia el Dockerfile en tu VM de TPU

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

5. Establece una conexión SSH a la VM de TPU

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

6. Compila la imagen de Docker

   sudo docker build -t your-image-name . 
   

7. Inicia el contenedor de Docker

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

8. Configura variables de entorno

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

Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.

1. Escribe exit para salir del contenedor de Docker.
2. Escribe exit para salir de la VM de TPU
3. Borra la VM de TPU

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

pod de TPU

1. Crea un archivo llamado Dockerfile en tu directorio actual y pega el texto siguiente

   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. Crea una VM de 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. Copia el Dockerfile en tu VM de TPU

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

4. Establece una conexión SSH a la VM de TPU

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

5. Compila la imagen de Docker

   sudo docker build -t your-image-name . 
   

6. Inicia un contenedor de Docker

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

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

Dispositivo de TPU

1. Crear una VM de Cloud TPU

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

2. Establece una conexión SSH a la VM de TPU

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

3. Inicia un contenedor en la VM de TPU con la imagen de PyTorch/XLA nocturna.

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

4. Configura el entorno de ejecución de la TPU

   Existen dos opciones de entorno de ejecución de PyTorch/XLA: PJRT y XRT. Te recomendamos puede usar PJRT a menos que tenga un motivo. Para obtener más información sobre el diferentes configuraciones de entorno de ejecución, verás que tienes un motivo para usar XRT. Para Para obtener más información sobre las diferentes opciones de configuración del entorno de ejecución, consulta la documentación sobre el entorno de ejecución de PJRT.

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

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

5. Clona el repositorio de XLA de PyTorch

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

6. Entrenar ResNet50

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

Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.

1. Escribe exit para salir del contenedor de Docker.
2. Escribe exit para salir de la VM de TPU
3. Borra la VM de TPU

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

pod de TPU

Cuando ejecutas código de PyTorch en un pod de TPU, debes hacerlo en todas las TPU trabajadores al mismo tiempo. Una forma de hacerlo es con la el comando gcloud compute tpus tpu-vm ssh con --worker=all y --command. En el siguiente procedimiento, se muestra cómo crear una instancia para facilitar la configuración de cada trabajador TPU.

1. Crea una VM de TPU

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

2. Agrega el usuario actual al grupo de Docker

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

3. Ejecuta la secuencia de comandos de entrenamiento en un contenedor en todos los trabajadores 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"
   

   Marcas de comando de Docker:

   • --rm quita el contenedor después de que finaliza su proceso.
   • --privileged expone el dispositivo de TPU al contenedor.
   • --net=host vincula todos los puertos del contenedor a la VM de TPU para permitir la comunicación entre los hosts en el Pod.
   • -e configuró variables de entorno.

Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.

Borra la VM de TPU con el siguiente comando:

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

Dispositivo de TPU

1. Crea la VM de TPU

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

2. Establece una conexión SSH a la VM de TPU

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

3. Iniciar el daemon de Docker en la VM de TPU

   sudo systemctl start docker
   

4. Inicia el contenedor de Docker

   sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
   

5. Instala JAX

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

6. Instalar FLAX

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

7. Ejecuta la secuencia de comandos de entrenamiento de 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
   

Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.

1. Escribe exit para salir del contenedor de Docker.
2. Escribe exit para salir de la VM de TPU

3. Borra la VM de TPU

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

pod de TPU

Cuando ejecutas el código JAX en un Pod de TPU, debes ejecutar tu código JAX en todas las TPU trabajadores al mismo tiempo. Una forma de hacerlo es usar gcloud compute tpus tpu-vm ssh con las marcas --worker=all y --command. Lo siguiente en el que se muestra cómo crear una imagen de Docker para configurar cada TPU a los trabajadores de una manera más fácil.

1. Crea un archivo llamado Dockerfile en tu directorio actual y pega el texto siguiente

   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. Compila la imagen de Docker

   docker build -t your-image-name .
   

3. Agrega una etiqueta a tu imagen de Docker antes de enviarla a Artifact Registry. Para obtener más información sobre cómo trabajar con Artifact Registry, consulta Trabaja con imágenes de contenedor.

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

4. Envía tu imagen de Docker a Artifact Registry

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

5. Crea una VM de TPU

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

6. Extrae la imagen de Docker de Artifact Registry en todos los trabajadores 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. Ejecuta el contenedor en todos los trabajadores de 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. Ejecuta la secuencia de comandos de entrenamiento en todos los trabajadores 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"
   

Cuando se complete la secuencia de comandos de entrenamiento, asegúrate de limpiar los recursos.

1. Cierra el contenedor en todos los trabajadores:

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

2. Borra la VM de TPU con el siguiente comando:

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