Executar cargas de trabalho de 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. É 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 compatível com o 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 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. Conecte-se à VM da TPU usando SSH

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

3. Verifique se o usuário Google Cloud recebeu a função de leitor do Artifact Registry. Para mais informações, consulte Conceder papéis do Artifact Registry.

4. Iniciar um contêiner na VM de TPU usando a imagem noturna do PyTorch/XLA

   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 de executar esse comando, o prompt de comando vai mudar, indicando que o terminal está conectado ao contêiner em execução.

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

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

   PJRT

   export PJRT_DEVICE=TPU

   XRT

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

6. Clone o repositório PyTorch XLA

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

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

Fração de TPU

Ao executar o código PyTorch em uma fração de TPU, é necessário executar o 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 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. Clone o repositório 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. Executar 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 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 define variáveis de ambiente.

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

Exclua a VM da TPU usando o comando a seguir:

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. Conecte-se à VM da TPU usando SSH

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

3. Iniciar o daemon do Docker na VM de TPU

   sudo systemctl start docker

4. Iniciar 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. Instalar o JAX

   pip install jax[tpu]

6. Instalar o FLAX

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

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

3. Excluir a VM da 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, é preciso executar o código JAX 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 flags --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.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. Preparar 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. Compilar a imagem Docker

   docker build -t your-image-name .

4. 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-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

5. 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

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

7. Extraia a imagem do Docker do Artifact Registry em todos os trabalhadores 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. Executar 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-name:your-tag bash"

9. Executar 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. Desativar 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. Excluir a VM da TPU

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

A seguir

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