Una vez que tu código de PyTorch se ejecute en una sola VM de TPU, puedes escalarlo de forma vertical en una porción de TPU.
Las porciones de TPU son varias placas de TPU conectadas entre sí en conexiones de red dedicadas de alta velocidad. Este documento es una introducción a la ejecución de código de PyTorch en porciones de TPU.
Crea una porción de Cloud TPU
Define algunas variables de entorno para facilitar el uso de los comandos.
El Google Cloud ID de tu proyecto. Usa un proyecto existente o
crea uno nuevo.
TPU_NAME
El nombre de la TPU.
ZONE
Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU.
ACCELERATOR_TYPE
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas
crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
Después de crear la porción de TPU, debes instalar PyTorch en todos los hosts en la porción de TPU. Puedes hacerlo con el comando gcloud compute tpus tpu-vm ssh usando
los parámetros --worker=all y --commamnd.
Si los siguientes comandos fallan debido a un error de conexión SSH, es posible que las VMs de TPU no tengan direcciones IP externas. Para acceder a una VM de TPU sin una dirección IP externa, sigue las instrucciones que se indican en Cómo conectarse a una VM de TPU sin una dirección IP pública.
Instala PyTorch/XLA en todos los trabajadores de VM de TPU:
Ejecuta una secuencia de comandos de entrenamiento en tu porción de TPU
Ejecuta la secuencia de comandos de entrenamiento en todos los trabajadores. La secuencia de comandos de entrenamiento usa una estrategia de fragmentación de datos múltiples de un solo programa (SPMD). Para obtener más información sobre SPMD, consulta la Guía del usuario de SPMD de PyTorch/XLA.
El entrenamiento dura alrededor de 15 minutos. Cuando se complete, deberías ver un mensaje similar al siguiente:
Epoch 1 test end 23:49:15, Accuracy=100.00
10.164.0.11 [0] Max Accuracy: 100.00%
Limpia
Cuando termines de usar la VM de TPU, sigue estos pasos para limpiar los recursos.
Desconéctate de la instancia de Cloud TPU, si aún no lo hiciste:
(vm)$exit
El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.
Borra tus recursos de Cloud TPU.
$gcloudcomputetpustpu-vmdelete\--zone=${ZONE}
Ejecuta gcloud compute tpus tpu-vm list para verificar que los recursos se hayan borrado. La eliminación puede tardar varios minutos. El resultado del siguiente comando no debe incluir ninguno de los recursos creados en este instructivo:
[[["Fácil de comprender","easyToUnderstand","thumb-up"],["Resolvió mi problema","solvedMyProblem","thumb-up"],["Otro","otherUp","thumb-up"]],[["Difícil de entender","hardToUnderstand","thumb-down"],["Información o código de muestra incorrectos","incorrectInformationOrSampleCode","thumb-down"],["Faltan la información o los ejemplos que necesito","missingTheInformationSamplesINeed","thumb-down"],["Problema de traducción","translationIssue","thumb-down"],["Otro","otherDown","thumb-down"]],["Última actualización: 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 ```"]]
