Execute código PyTorch em fatias de TPU

Antes de executar os comandos neste documento, certifique-se de que seguiu as instruções em Configure uma conta e um projeto do Cloud TPU.

Depois de ter o código PyTorch em execução numa única VM da TPU, pode dimensionar o código executando-o numa divisão da TPU. As fatias de TPU são várias placas de TPU ligadas entre si através de ligações de rede de alta velocidade dedicadas. Este documento é uma introdução à execução de código PyTorch em fatias de TPU.

Crie uma fatia da Cloud TPU

1. Defina algumas variáveis de ambiente para facilitar a utilização dos comandos.

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

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
   RUNTIME_VERSION	A versão do software do 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}

Instale o PyTorch/XLA na sua fatia

Depois de criar a fatia de TPU, tem de instalar o PyTorch em todos os anfitriões na fatia de TPU. Pode fazê-lo através do comando gcloud compute tpus tpu-vm ssh com os parâmetros --worker=all e --commamnd.

Se os seguintes comandos falharem devido a um erro de ligação SSH, pode dever-se ao facto de as VMs de TPU não terem endereços IP externos. Para aceder a uma VM da TPU sem um endereço IP externo, siga as instruções em Estabeleça ligação a uma VM da TPU sem um endereço IP público.

1. Instale o PyTorch/XLA em todos os trabalhadores da VM da 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. Clonar o XLA em todos os trabalhadores da 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"

Execute um script de preparação na sua fatia de TPU

Execute o script de preparação em todos os trabalhadores. O script de preparação usa uma estratégia de divisão em partições de programa único e vários dados (SPMD). Para mais informações sobre SPMD, consulte o guia do utilizador de 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"

A preparação demora cerca de 15 minutos. Quando estiver concluído, deve ver uma mensagem semelhante à seguinte:

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 estes passos para limpar os recursos.

1. Desligue-se da instância do Cloud TPU, se ainda não o tiver feito:

   (vm)$ exit

   O comando deve agora ser username@projectname, o que indica que está no Cloud Shell.

2. Elimine os seus recursos do Cloud TPU.

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

3. Verifique se os recursos foram eliminados executando gcloud compute tpus tpu-vm list. A eliminação pode demorar vários minutos. O resultado do seguinte comando não deve incluir nenhum dos recursos criados neste tutorial:

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