Execute cargas de trabalho de TPU num contentor do Docker

Os contentores Docker facilitam a configuração das aplicações combinando o seu código e todas as dependências necessárias num pacote distribuível. Pode executar contentores Docker em VMs de TPU para simplificar a configuração e a partilha das suas aplicações Cloud TPU. Este documento descreve como configurar um contentor Docker para cada framework de ML suportado pela Cloud TPU.

Importante: pode executar um máximo de 256 pods do Docker numa VM de TPU.

Prepare um modelo do PyTorch num contentor do Docker

Dispositivo 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. Estabeleça ligação à VM de TPU através de SSH

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

3. Certifique-se de que ao seu Google Cloud utilizador foi concedido o papel de leitor do Artifact Registry. Para mais informações, consulte o artigo Conceder funções do Artifact Registry.

4. Inicie um contentor na VM da TPU com 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

   Nota: depois de executar este comando, a linha de comandos muda para indicar que o terminal está ligado ao contentor em execução.

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

   Existem duas opções de tempo de execução do PyTorch/XLA: PJRT e XRT. Recomendamos que use o PJRT, a menos que tenha um motivo para usar o XRT. Para saber mais acerca das diferentes configurações de tempo de execução, consulte a documentação de tempo 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. Preparar ResNet50

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

Quando o script de preparação estiver concluído, certifique-se de que limpa os recursos.

1. Escreva exit para sair do contentor Docker
2. Escreva exit para sair da VM da TPU

3. Elimine a VM da TPU

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

Fatia de TPU

Quando executa código PyTorch numa fatia de TPU, tem de executar o código em todos os trabalhadores de TPU ao mesmo tempo. Uma forma de o fazer é usar o comando gcloud compute tpus tpu-vm ssh com os indicadores --worker=all e --command. O procedimento seguinte mostra como criar uma imagem do Docker para facilitar a configuração de cada worker de TPU.

1. Crie 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. Adicione o utilizador atual ao grupo 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 preparação num contentor em todos os trabalhadores da 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"

   Sinalizadores de comandos do Docker:

   • --rm remove o contentor após a terminação do respetivo processo.
   • --privileged expõe o dispositivo TPU ao contentor.
   • --net=host associa todas as portas do contentor à VM da TPU para permitir a comunicação entre os anfitriões no pod.
   • -e define variáveis de ambiente.

Quando o script de preparação estiver concluído, certifique-se de que limpa os recursos.

Elimine a VM da TPU com o seguinte comando:

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

Prepare um modelo JAX num contentor do Docker

Dispositivo TPU

1. Crie a VM da TPU

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

2. Estabeleça ligação à VM de TPU através de SSH

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

3. Inicie o daemon Docker na VM da TPU

   sudo systemctl start docker

4. Inicie o contentor 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 pacotes tensorflow e tensorflow-dataset

   pip install tensorflow
   pip install tensorflow-datasets

8. Execute o script de preparação do FLAX MNIST

   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 preparação estiver concluído, certifique-se de que limpa os recursos.

1. Escreva exit para sair do contentor Docker
2. Escreva exit para sair da VM da TPU

3. Elimine a VM da TPU

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

Fatia de TPU

Quando executa código JAX numa fatia de TPU, tem de executar o código JAX em todos os trabalhadores de TPU ao mesmo tempo. Uma forma de o fazer é usar o comando gcloud compute tpus tpu-vm ssh com as flags --worker=all e --command. O procedimento seguinte mostra como criar uma imagem do Docker para facilitar a configuração de cada worker da TPU.

1. Crie um ficheiro denominado Dockerfile no seu diretório atual e cole o seguinte texto

   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. Crie a imagem de Docker

   docker build -t your-image-name .

4. Adicione uma etiqueta à sua imagem do Docker antes de a enviar para o Artifact Registry. Para mais informações sobre como trabalhar com o Artifact Registry, consulte o artigo Trabalhe com imagens de contentores.

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

5. Envie a sua imagem de 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 da 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 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"

8. Executar o contentor em todos os trabalhadores 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. Execute o script de preparação em todos os trabalhadores 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 preparação estiver concluído, certifique-se de que limpa os recursos.

1. Encerre o contentor em todos os trabalhadores

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

2. Elimine a VM da TPU

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

O que se segue?

• Tutoriais do Cloud TPU
• Faça a gestão das TPUs
• Arquitetura do Cloud TPU
• Execute código JAX em fatias de TPUs