Executar cargas de trabalho de TPU em um contêiner do Docker

Os contêineres do Docker facilitam a configuração de aplicativos com a combinação do seu código e de 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 documento descreve como configurar um contêiner do Docker para cada framework de ML aceito pelo Cloud TPU.

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

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

Dispositivo de TPU

1. Crie 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. Conecte-se à VM usando SSH.

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

3. Verifique se o usuário do Google Cloud recebeu o papel de Leitor do Artifact Registry. Para mais informações, consulte Como conceder papéis do Artifact Registry.

4. Inicie um contêiner na VM de TPU usando a imagem do PyTorch/XLA criada durante a noite.

   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

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

5. Configure o ambiente de execução da TPU.

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

6. Clone o repositório do PyTorch/XLA.

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

7. Treine 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 de TPU.

3. Exclua a VM de TPU.

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

Fração de TPU

Ao executar o código PyTorch em uma fração de TPU, isso deve ser feito em todos os workers de TPU ao mesmo tempo. Uma forma 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 Docker para facilitar a configuração de cada worker de TPU.

1. Crie 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. Adicione 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. Clone o repositório do 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. Execute o script de treinamento em um contêiner em todos os workers 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"

   Flags de comando do Docker:

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

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

Exclua a VM de TPU executando este 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 de TPU

1. Crie 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. Conecte-se à VM usando SSH.

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

3. Inicie o daemon do Docker na VM de TPU.

   sudo systemctl start docker

4. Inicie o contêiner do 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. Instale o JAX.

   pip install jax[tpu]

6. Instale o FLAX.

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

7. Instale os pacotes tensorflow e tensorflow-dataset.

   pip install tensorflow
   pip install tensorflow-datasets

8. 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 de TPU.

3. Exclua a VM de TPU.

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

Fração de TPU

Ao executar o código JAX em uma fração de TPU, isso deve ser feito em todos os workers de TPU ao mesmo tempo. Uma forma 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 Docker para facilitar a configuração de cada worker de TPU.

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

   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. Prepare um 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. Compile a imagem Docker.

   docker build -t your-image-name .

4. Adicione uma tag à imagem Docker antes de enviá-la ao Artifact Registry. Para saber como trabalhar com o Artifact Registry, consulte Trabalhar com imagens de contêiner.

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

5. Envie a imagem Docker para o Artifact Registry.

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

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

7. Extraia a imagem Docker do Artifact Registry em todos os workers 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. Execute o contêiner em todos os workers 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-name:your-tag bash"

9. Execute o script de treinamento em todos os workers 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"

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

1. Desative 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 de TPU.

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

A seguir

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