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
Defina algumas variáveis de ambiente para facilitar o uso dos comandos.
O ID do seu Google Cloud projeto. Use um projeto existente ou
crie um novo.
TPU_NAME
O nome da TPU.
ZONE
A zona em que a VM TPU será criada. Para mais informações sobre as zonas com suporte, consulte
Regiões e zonas de TPU.
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 com cada versão de TPU, consulte
Versões de TPU.
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 externo. 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.
Instale o PyTorch/XLA em todos os workers de VM de TPU:
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.
O treinamento leva cerca de 15 minutos. Quando ela 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.
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.
Exclua seus recursos do Cloud TPU.
$gcloudcomputetpustpu-vmdelete\--zone=${ZONE}
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:
[[["Fácil de entender","easyToUnderstand","thumb-up"],["Meu problema foi resolvido","solvedMyProblem","thumb-up"],["Outro","otherUp","thumb-up"]],[["Difícil de entender","hardToUnderstand","thumb-down"],["Informações incorretas ou exemplo de código","incorrectInformationOrSampleCode","thumb-down"],["Não contém as informações/amostras de que eu preciso","missingTheInformationSamplesINeed","thumb-down"],["Problema na tradução","translationIssue","thumb-down"],["Outro","otherDown","thumb-down"]],["Última atualização 2025-09-04 UTC."],[],[],null,["# Run PyTorch code on TPU slices\n==============================\n\nBefore running the commands in this document, make sure you have followed the\ninstructions in [Set up an account and Cloud TPU project](/tpu/docs/setup-gcp-account).\n\nAfter you have your PyTorch code running on a single TPU VM, you can scale up\nyour code by running it on a [TPU slice](/tpu/docs/system-architecture-tpu-vm#slices).\nTPU slices are multiple TPU boards connected to each other over dedicated\nhigh-speed network connections. This document is an introduction to running\nPyTorch code on TPU slices.\n\nCreate a Cloud TPU slice\n------------------------\n\n1. Define some environment variables to make the commands easier to use.\n\n\n ```bash\n export PROJECT_ID=your-project-id\n export TPU_NAME=your-tpu-name\n export ZONE=europe-west4-b\n export ACCELERATOR_TYPE=v5p-32\n export RUNTIME_VERSION=v2-alpha-tpuv5\n ``` \n\n #### Environment variable descriptions\n\n \u003cbr /\u003e\n\n2. Create your TPU VM by running the following command:\n\n ```bash\n $ gcloud compute tpus tpu-vm create ${TPU_NAME} \\\n --zone=${ZONE} \\\n --project=${PROJECT_ID} \\\n --accelerator-type=${ACCELERATOR_TYPE} \\\n --version=${RUNTIME_VERSION}\n ```\n\nInstall PyTorch/XLA on your slice\n---------------------------------\n\nAfter creating the TPU slice, you must install PyTorch on all hosts in the\nTPU slice. You can do this using the `gcloud compute tpus tpu-vm ssh` command using\nthe `--worker=all` and `--commamnd` parameters.\n\nIf the following commands fail due to an SSH connection error, it might be\nbecause the TPU VMs don't have external IP addresses. To access a TPU VM without\nan external IP address, follow the instructions in [Connect to a TPU VM without\na public IP address](/tpu/docs/tpu-iap).\n\n1. Install PyTorch/XLA on all TPU VM workers:\n\n ```bash\n gcloud compute tpus tpu-vm ssh ${TPU_NAME} \\\n --zone=${ZONE} \\\n --project=${PROJECT_ID} \\\n --worker=all \\\n --command=\"pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html\"\n ```\n2. Clone XLA on all TPU VM workers:\n\n ```bash\n gcloud compute tpus tpu-vm ssh ${TPU_NAME} \\\n --zone=${ZONE} \\\n --project=${PROJECT_ID} \\\n --worker=all \\\n --command=\"git clone https://github.com/pytorch/xla.git\"\n ```\n\nRun a training script on your TPU slice\n---------------------------------------\n\nRun the training script on all workers. The training script uses a Single Program\nMultiple Data (SPMD) sharding strategy. For more information on SPMD, see\n[PyTorch/XLA SPMD User Guide](https://pytorch.org/xla/release/r2.4/spmd.html). \n\n```bash\ngcloud compute tpus tpu-vm ssh ${TPU_NAME} \\\n --zone=${ZONE} \\\n --project=${PROJECT_ID} \\\n --worker=all \\\n --command=\"PJRT_DEVICE=TPU python3 ~/xla/test/spmd/test_train_spmd_imagenet.py \\\n --fake_data \\\n --model=resnet50 \\\n --num_epochs=1 2\u003e&1 | tee ~/logs.txt\"\n```\n\nThe training takes about 15 minutes. When it completes, you should see a message\nsimilar to the following: \n\n```\nEpoch 1 test end 23:49:15, Accuracy=100.00\n 10.164.0.11 [0] Max Accuracy: 100.00%\n```\n\nClean up\n--------\n\nWhen you are done with your TPU VM, follow these steps to clean up your resources.\n\n1. Disconnect from the Cloud TPU instance, if you have not already\n done so:\n\n ```bash\n (vm)$ exit\n ```\n\n Your prompt should now be `username@projectname`, showing you are in the\n Cloud Shell.\n2. Delete your Cloud TPU resources.\n\n ```bash\n $ gcloud compute tpus tpu-vm delete \\\n --zone=${ZONE}\n ```\n3. Verify the resources have been deleted by running `gcloud compute tpus tpu-vm list`. The\n deletion might take several minutes. The output from the following command\n shouldn't include any of the resources created in this tutorial:\n\n ```bash\n $ gcloud compute tpus tpu-vm list --zone=${ZONE}\n ```"]]
