Execute aplicativos da Cloud TPU em um contêiner do Docker

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

Importante: é possível executar no máximo 256 pods do Docker em uma VM da 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. criar o bucket do Cloud Storage

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

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.16.1-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. 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, 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 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 de TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.16.1-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. 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 da 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 PyTorch/XLA noturna.

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

Quando você executa o código JAX em um pod de TPU, é necessário executar esse código em todos os workers de 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. 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 de 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 frações do Pod de TPU
• Executar o código JAX em frações do Pod de TPU