Executar aplicativos do Cloud TPU em um contêiner do Docker

Os contêineres do Docker facilitam a configuração de aplicativos ao combinar seu código e todas as dependências necessárias em um pacote distribuível. É possível executar contêineres do Docker em VMs de TPU para simplificar a configuração e o compartilhamento de aplicativos do Cloud TPU. Este document descreve como configurar um contêiner do Docker para cada framework de ML com suporte ao Cloud TPU.

Importante: é possível executar no máximo 256 pods do Docker em uma VM de TPU.

Treinar um modelo do TensorFlow em um contêiner do Docker

Dispositivo TPU

1. Crie um arquivo com o nome Dockerfile no diretório atual e cole o texto seguinte

   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. crie o bucket do Cloud Storage

   gcloud storage buckets create gs://your-bucket-name --location=europe-west4

3. Criar uma 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. Copie o Dockerfile para a VM da TPU

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

5. Conectar-se por SSH à VM de TPU

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

6. Compilar a imagem Docker

   sudo docker build -t your-image-name . 

7. Inicie o contêiner do Docker

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

8. Defina as variáveis de ambiente

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

Quando o script de treinamento for concluído, limpe os recursos.

1. Digite exit para sair do contêiner do Docker.
2. Digite exit para sair da VM da TPU
3. Excluir a VM da TPU

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

Pod de TPU

1. Crie um arquivo com o nome Dockerfile no diretório atual e cole o texto seguinte

   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. Criar uma 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. Copie o Dockerfile para a VM da TPU

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

4. Conectar-se por SSH à VM de TPU

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

5. Compilar a imagem Docker

   sudo docker build -t your-image-name . 

6. Iniciar um contêiner do Docker

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

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

Quando o script de treinamento for concluído, limpe os recursos.

1. Digite exit para sair do contêiner do Docker.
2. Digite exit para sair da VM da TPU
3. Exclua a VM da TPU

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

Treinar um modelo PyTorch em um contêiner do Docker

Dispositivo TPU

1. Criar uma VM do Cloud TPU

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

2. SSH na VM da TPU

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

3. Inicie um contêiner na VM da TPU usando a imagem noturna PyTorch/XLA.

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

   Observação: depois de executar este comando, o prompt de comando vai mudar, indicando que o terminal está conectado ao contêiner em execução.

4. Configurar o ambiente de execução da TPU

   Há duas opções de execução do PyTorch/XLA: PJRT e XRT. Recomendamos que você usar o PJRT, a menos que você tenha um motivo para usar o XRT. Para saber mais sobre o diferentes configurações de tempo de execução, então há um motivo para usar o XRT. Para Para saber mais sobre as diferentes configurações do ambiente de execução, consulte a documentação do ambiente de execução da PJRT.

   PJRT

   export PJRT_DEVICE=TPU

   XRT

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

5. Clonar o repositório PyTorch XLA

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

6. Treinar ResNet50

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

Quando o script de treinamento for concluído, limpe os recursos.

1. Digite exit para sair do contêiner do Docker.
2. Digite exit para sair da VM da TPU
3. Excluir a VM da TPU

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

Pod de TPU

Ao executar o código PyTorch em um pod de TPU, é necessário executar o código em todos os workers do TPU ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as flags --worker=all e --command. O procedimento a seguir mostra como criar uma imagem do Docker para facilitar a configuração de cada worker de TPU.

1. Criar uma 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. Adicionar o usuário atual ao grupo do Docker

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

3. Execute o script de treinamento em um contêiner em todos os workers da 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"

   Sinalizações de comando do Docker:

   • --rm remove o contêiner após o encerramento do processo.
   • --privileged expõe o dispositivo TPU ao contêiner.
   • --net=host vincula todas as portas do contêiner à VM da TPU para permitir a comunicação entre os hosts no pod.
   • -e definiu as variáveis de ambiente.

Quando o script de treinamento for concluído, limpe os recursos.

Exclua a VM da TPU usando o seguinte comando:

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

Treinar um modelo JAX em um contêiner do Docker

Dispositivo TPU

1. Criar a 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. SSH na VM de TPU

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

3. Iniciar o daemon do Docker na VM da TPU

   sudo systemctl start docker

4. Iniciar contêiner do Docker

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

5. Instalar o JAX

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

6. Instalar o FLAX

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

7. Execute o script de treinamento MNIST do FLAX

   cd flax/examples/mnist
   python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5

Quando o script de treinamento for concluído, limpe os recursos.

1. Digite exit para sair do contêiner do Docker.
2. Digite exit para sair da VM da TPU

3. Excluir a VM da TPU

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

Pod de TPU

Ao executar o código JAX em um pod de TPU, é necessário executar esse código em todas as TPUs ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as flags --worker=all e --command. O seguinte mostra como criar uma imagem Docker para configurar cada TPU profissional mais fácil.

1. Crie um arquivo chamado Dockerfile no diretório atual e cole o seguinte 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. Compilar a imagem Docker

   docker build -t your-image-name .

3. Adicione uma tag à imagem do Docker antes de enviá-la para o Artifact Registry. Para mais informações sobre como trabalhar com o Artifact Registry, consulte Trabalhar com imagens de contêiner.

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

4. Envie a imagem do Docker para o Artifact Registry

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

5. Criar uma VM da TPU

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

6. Extraia a imagem do Docker do Artifact Registry em todos os workers da 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. Execute o contêiner em todos os workers da 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. Execute o script de treinamento em todos os workers da 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"

Quando o script de treinamento for concluído, limpe os recursos.

1. Encerre o contêiner em todos os workers:

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

2. Exclua a VM da TPU usando o seguinte comando:

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

A seguir

• Tutoriais da Cloud TPU
• Gerenciar TPUs
• Arquitetura do sistema do Cloud TPU
• Executar o código do TensorFlow em Pod de TPU TPU
• Executar o código JAX em frações do pod de TPU