Ejecuta cargas de trabajo de TPU en un contenedor de Docker

Los contenedores de Docker facilitan la configuración de las aplicaciones, ya que combinan tu código y todas las dependencias necesarias en un paquete distribuible. Puedes ejecutar contenedores de Docker dentro de VMs de TPU para simplificar la configuración y el uso compartido de tus aplicaciones de Cloud TPU. En este documento, se describe cómo configurar un contenedor de Docker para cada framework de AA compatible con Cloud TPU.

Importante: Puedes ejecutar un máximo de 256 Pods de Docker en una VM de TPU.

Entrena un modelo de PyTorch en un contenedor de Docker

Dispositivo de TPU

1. Crea 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. Conéctate a la VM de TPU con SSH

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

3. Asegúrate de que a tu usuario Google Cloud se le haya otorgado el rol de lector de Artifact Registry. Para obtener más información, consulta Cómo otorgar roles de Artifact Registry.

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

   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

   Nota: Después de ejecutar este comando, el símbolo del sistema cambiará para indicar que la terminal está conectada al contenedor en ejecución.

5. Configura el tiempo de ejecución de la TPU

   Hay dos opciones de tiempo de ejecución de PyTorch/XLA: PJRT y XRT. Te recomendamos que uses PJRT, a menos que tengas un motivo para usar XRT. Para obtener más información sobre las diferentes configuraciones de tiempo de ejecución, consulta la documentación del tiempo de ejecución de PJRT.

   PJRT

   export PJRT_DEVICE=TPU

   XRT

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

6. Clona el repositorio de PyTorch XLA

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

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

Porción de TPU

Cuando ejecutas código PyTorch en una porción de TPU, debes ejecutar el código en todos los trabajadores de TPU al mismo tiempo. Una forma de hacerlo es usar el comando gcloud compute tpus tpu-vm ssh con las marcas --worker=all y --command. En el siguiente procedimiento, se muestra cómo crear una imagen de Docker 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. Clona el repositorio de 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. Ejecuta la secuencia de comandos de entrenamiento en un contenedor en todos los trabajadores de 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"

   Marcas de comandos 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 del pod.
   • -e establece 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

Entrena un modelo de JAX en un contenedor de Docker

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. Conéctate a la VM de TPU con SSH

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

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

   sudo systemctl start docker

4. Inicia el contenedor de 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. Instala JAX

   pip install jax[tpu]

6. Instala FLAX

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

7. Instala los paquetes tensorflow y tensorflow-dataset

   pip install tensorflow
   pip install tensorflow-datasets

8. Ejecuta la secuencia de comandos de entrenamiento de MNIST de FLAX

   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

Porción de TPU

Cuando ejecutas código JAX en una porción de TPU, debes ejecutar el código JAX en todos los trabajadores de TPU al mismo tiempo. Una forma de hacerlo es usar el comando gcloud compute tpus tpu-vm ssh con las marcas --worker=all y --command. En el siguiente procedimiento, se muestra cómo crear una imagen de Docker para facilitar la configuración de cada trabajador de TPU.

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

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

   docker build -t your-image-name .

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

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

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

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

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

¿Qué sigue?

• Instructivos de Cloud TPU
• Administra TPU
• Arquitectura de Cloud TPU
• Ejecuta el código JAX en porciones de TPU