Executar o código PyTorch em fatias 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 de TPU. As frações 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 fatias de TPU.

Criar uma fração 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 sua fatia

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

Se os comandos a seguir falharem devido a um erro de conexão SSH, talvez seja porque as VMs TPU não têm endereços IP externos. Para acessar uma VM da TPU sem um endereço IP externo, siga as instruções em Conectar-se a uma VM da TPU sem um endereço IP público.

1. Instale 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 da 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 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 15 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 Cloud TPU, 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 seus recursos do Cloud TPU.

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

3. Execute gcloud compute tpus tpu-vm 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}