Ejecuta aplicaciones de Cloud TPU en un contenedor de Docker

Los contenedores de Docker facilitan la configuración de 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 que admite Cloud TPU.

Entrena un modelo de TensorFlow en un contenedor de Docker

Dispositivo de TPU

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

    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

    gcloud storage buckets create gs://your-bucket-name --location=europe-west4
  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.18.0-pjrt
  4. Cómo copiar el Dockerfile en tu VM de TPU

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
  5. Conéctate a la VM de TPU mediante SSH

    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. Configure las 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 siguiente texto:

    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.18.0-pod-pjrt
  3. Cómo copiar el Dockerfile en tu VM de TPU

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
  4. Conéctate a la VM de TPU mediante SSH

    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"

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

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 a través de SSH

    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 nocturna de PyTorch/XLA.

    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 TPU

    Hay dos opciones de entorno 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 del entorno de ejecución, consulta si tienes motivos para usar XRT. Para obtener más información sobre las diferentes configuraciones del entorno de ejecución, consulta la documentación del 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 PyTorch XLA

    git clone --recursive https://github.com/pytorch/xla.git
  6. 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

pod de TPU

Cuando ejecutas código PyTorch en un pod de TPU, debes ejecutarlo 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 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 de 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 del 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 del pod.
    • -e establece las 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 a través de 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 -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. Instala 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 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

pod de TPU

Cuando ejecutas código JAX en un pod de TPU, debes ejecutarlo 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.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 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"
  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 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. Apaga 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

Entrena un modelo de JAX en un contenedor de Docker con la pila estable de JAX

Puedes compilar las imágenes de Docker de MaxText y MaxDiffusion con la imagen base de JAX Stable Stack.

El JAX Stable Stack proporciona un entorno coherente para MaxText y MaxDiffusion, ya que agrupa JAX con paquetes principales, como orbax, flax, optax y libtpu.so. Estas bibliotecas se prueban para garantizar la compatibilidad y proporcionar una base estable para compilar y ejecutar MaxText y MaxDiffusion. Esto elimina los posibles conflictos debido a versiones de paquetes incompatibles.

La pila estable de JAX incluye un libtpu.so completamente lanzado y calificado, la biblioteca principal que impulsa la compilación, la ejecución y la configuración de la red de ICI del programa de TPU. La versión de libtpu reemplaza la compilación nocturna que usaba JAX anteriormente y garantiza una funcionalidad coherente de los cálculos de XLA en TPU con pruebas de calificación a nivel de PJRT en IR de HLO/StableHLO.

Para compilar la imagen de Docker de MaxText y MaxDiffusion con la pila estable de JAX, cuando ejecutes la secuencia de comandos docker_build_dependency_image.sh, establece la variable MODE en stable_stack y la variable BASEIMAGE en la imagen base que deseas usar.

En el siguiente ejemplo, se especifica us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1 como la imagen 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
   

Para obtener una lista de las imágenes base de JAX Stable Stack disponibles, consulta Imágenes de JAX Stable Stack en Artifact Registry.

¿Qué sigue?