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

Antes de executar os comandos neste documento, verifique se você seguiu as instruções em Configurar uma conta e um projeto do Cloud TPU.

Depois de executar o código PyTorch em uma única VM de TPU, é possível escalonar verticalmente o código executando-o em uma fração do pod de TPU. As frações do pod de TPU são várias placas de TPU conectadas entre si por conexões de rede dedicadas de alta velocidade. Este documento é uma introdução à execução do código PyTorch em frações do pod de TPU.

Criar uma fração de pod da Cloud TPU

1. Defina algumas variáveis de ambiente para facilitar o uso dos comandos.

   export PROJECT_ID=your-project
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=europe-west4-b
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export TPU_NAME=your-tpu-name

   Descrições das variáveis de ambiente

   PROJECT_ID

   O ID do Google Cloud projeto.

   ACCELERATOR_TYPE

   O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis para cada versão de TPU, consulte Versões de TPU.

   ZONE

   A zona em que você planeja criar a Cloud TPU.

   RUNTIME_VERSION

   A versão do software da Cloud TPU.

   TPU_NAME

   O nome atribuído pelo usuário ao Cloud TPU.

2. Crie a VM da TPU executando o seguinte comando:

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

Instalar o PyTorch/XLA na fração do pod

Depois de criar a fração do pod de TPU, é necessário instalar o PyTorch em todos os hosts dessa fração. Para isso, use o comando gcloud compute tpus tpu-vm ssh com os parâmetros --worker=all e --commamnd.

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

2. 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 https://github.com/pytorch/xla.git"

Executar um script de treinamento na fração do pod da TPU

Execute o script de treinamento em todos os workers. O script de treinamento usa uma estratégia de fragmentação de dados de vários dados de um único programa (SPMD, na sigla em inglês). Para mais informações sobre SPMD, consulte o Guia do usuário do SPMD do PyTorch/XLA.

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

O treinamento leva cerca de cinco minutos. Quando ela for concluída, você vai receber 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}