Executar o código PyTorch em frações do pod de TPU

O PyTorch/XLA exige que todas as VMs de TPU tenham acesso ao código e aos dados do modelo. É possível usar um script de inicialização para fazer o download do software necessário para distribuir os dados do modelo a todas as VMs de TPU.

Se você estiver conectando as VMs de TPU a uma nuvem privada virtual (VPC), adicione uma regra de firewall ao projeto para permitir a entrada nas portas 8470 a 8479. Para mais informações sobre a adição de regras de firewall, consulte Como usar regras de firewall.

Configurar o ambiente

1. No Cloud Shell, execute o seguinte comando para garantir que você está executando a versão atual do gcloud:

   $ gcloud components update

   Se você precisar instalar o gcloud, use o seguinte comando:

   $ sudo apt install -y google-cloud-sdk

2. Crie algumas variáveis de ambiente:

   $ export PROJECT_ID=project-id
   $ export TPU_NAME=tpu-name
   $ export ZONE=us-central2-b
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ export ACCELERATOR_TYPE=v4-32

Criar a VM de TPU

$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone=${ZONE} \
--project=${PROJECT_ID} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version ${RUNTIME_VERSION}

Configurar e executar o script de treinamento

1. Adicione o certificado SSH ao seu projeto:

   ssh-add ~/.ssh/google_compute_engine

2. Instalar o PyTorch/XLA em todos os workers de VM de TPU

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all --command="
     pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

3. Clone o XLA em todos os workers de VM do TPU

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all --command="git clone -b r2.5 https://github.com/pytorch/xla.git"

4. Executar o script de treinamento em todos os workers

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all \
     --command="PJRT_DEVICE=TPU python3 ~/xla/test/test_train_mp_imagenet.py  \
     --fake_data \
     --model=resnet50  \
     --num_epochs=1 2>&1 | tee ~/logs.txt"
     

   O treinamento leva cerca de cinco minutos. Quando ele for concluída, você verá uma mensagem semelhante a esta:

   Epoch 1 test end 23:49:15, Accuracy=100.00
   10.164.0.11 [0] Max Accuracy: 100.00%
   

Limpar

Quando terminar de usar a VM de TPU, siga estas etapas para limpar os recursos.

1. Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:

   (vm)$ exit

   Agora, o prompt precisa ser username@projectname, mostrando que você está no Cloud Shell.

2. Exclua os recursos do Cloud TPU e do Compute Engine.

   $ gcloud compute tpus tpu-vm delete  \
     --zone=${ZONE}

3. Execute gcloud compute tpus execution-groups list para verificar se os recursos foram excluídos. A exclusão pode levar vários minutos. A saída do comando a seguir não pode incluir nenhum dos recursos criados neste tutorial:

   $ gcloud compute tpus tpu-vm list --zone=${ZONE}