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

Os contêineres Docker facilitam a configuração de aplicativos combinando seu código e todas as dependências necessárias em um pacote distribuível. É possível executar contêineres do Docker em VMs da TPU para simplificar a configuração e o compartilhamento dos aplicativos da Cloud TPU. Neste documento, descrevemos como configurar um contêiner do Docker para cada framework de ML compatível com o Cloud TPU.

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

Dispositivo TPU

1. Crie um arquivo chamado Dockerfile no diretório atual e cole o texto a seguir:

   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. Criação do bucket do Cloud Storage

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

3. Criar uma VM da TPU

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

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. Criar a imagem do Docker

   sudo docker build -t your-image-name .
   

7. Iniciar 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, certifique-se de limpar 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 chamado Dockerfile no diretório atual e cole o texto a seguir:

   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 da TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.13.0-pod
   

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. Criar a imagem do 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, certifique-se de limpar 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
   

Treinar um modelo PyTorch em um contêiner do Docker

Dispositivo TPU

1. Criar 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 do 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 esse comando, o prompt de comando 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 ambiente de execução do PyTorch/XLA: PJRT e XRT. Recomendamos que você use o PJRT, a menos que tenha um motivo. Para saber mais sobre as diferentes configurações de ambiente de execução, veja por que você precisa usar o XRT. Para saber mais sobre as diferentes configurações do ambiente de execução, consulte a documentação do ambiente de execução do PJRT.

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

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

5. Clone 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, certifique-se de limpar 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, é preciso executar o código em todos os workers da TPU ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as sinalizações --worker=all e --command. O procedimento a seguir mostra como criar uma imagem do Docker para facilitar a configuração de cada worker da TPU.

1. Criar uma VM da 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 término do processo.
   • --privileged expõe o dispositivo da 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 variáveis de ambiente.

Quando o script de treinamento for concluído, certifique-se de limpar 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 da 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 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, certifique-se de limpar 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 executá-lo em todos os workers da TPU ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as sinalizações --worker=all e --command. O procedimento a seguir mostra como criar uma imagem do Docker para facilitar a configuração de cada worker da TPU.

1. Crie um arquivo chamado Dockerfile no diretório atual e cole o texto a seguir:

   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. Criar a imagem do Docker

   docker build -t your-image-name .
   

3. Adicione uma tag à imagem do Docker antes de enviá-la ao 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. Enviar 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, certifique-se de limpar 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 Pod de TPU TPU