Execute cargas de trabalho de TPU 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

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

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