Dopo aver eseguito il codice PyTorch su una singola VM TPU, puoi eseguire lo scaling up del codice eseguendolo su uno slice TPU.
Le sezioni TPU sono più schede TPU collegate tra loro tramite connessioni di rete ad alta velocità dedicate. Questo documento è un'introduzione all'esecuzione del codice PyTorch sulle sezioni TPU.
Creare un segmento Cloud TPU
Definisci alcune variabili di ambiente per semplificare l'utilizzo dei comandi.
Il tuo Google Cloud ID progetto. Utilizza un progetto esistente o
creane uno nuovo.
TPU_NAME
Il nome della TPU.
ZONE
La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta
Regioni e zone TPU.
ACCELERATOR_TYPE
Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta
Versioni TPU.
Dopo aver creato la sezione TPU, devi installare PyTorch su tutti gli host della sezione TPU. Puoi farlo utilizzando il comando gcloud compute tpus tpu-vm ssh con i parametri --worker=all e --commamnd.
Se i seguenti comandi non vanno a buon fine a causa di un errore di connessione SSH, potrebbe essere perché le VM TPU non hanno indirizzi IP esterni. Per accedere a una VM TPU senza un indirizzo IP esterno, segui le istruzioni riportate in Connettersi a una VM TPU senza un indirizzo IP pubblico.
Esegui uno script di addestramento sul tuo slice TPU
Esegui lo script di addestramento su tutti i worker. Lo script di addestramento utilizza una strategia di suddivisione in parti (SPMD) (Single Program
Multiple Data). Per ulteriori informazioni su SPMD, consulta la Guida dell'utente di PyTorch/XLA SPMD.
L'addestramento richiede circa 15 minuti. Al termine, dovresti visualizzare un messaggio simile al seguente:
Epoch 1 test end 23:49:15, Accuracy=100.00
10.164.0.11 [0] Max Accuracy: 100.00%
Esegui la pulizia
Al termine dell'utilizzo della VM TPU, segui questi passaggi per ripulire le risorse.
Se non l'hai già fatto, scollega l'istanza Cloud TPU:
(vm)$exit
Il tuo prompt dovrebbe ora essere username@projectname, a indicare che ti trovi in Cloud Shell.
Elimina le risorse Cloud TPU.
$gcloudcomputetpustpu-vmdelete\--zone=${ZONE}
Verifica che le risorse siano state eliminate eseguendo gcloud compute tpus tpu-vm list. L'eliminazione potrebbe richiedere alcuni minuti. L'output del seguente comando
non deve includere nessuna delle risorse create in questo tutorial:
[[["Facile da capire","easyToUnderstand","thumb-up"],["Il problema è stato risolto","solvedMyProblem","thumb-up"],["Altra","otherUp","thumb-up"]],[["Difficile da capire","hardToUnderstand","thumb-down"],["Informazioni o codice di esempio errati","incorrectInformationOrSampleCode","thumb-down"],["Mancano le informazioni o gli esempi di cui ho bisogno","missingTheInformationSamplesINeed","thumb-down"],["Problema di traduzione","translationIssue","thumb-down"],["Altra","otherDown","thumb-down"]],["Ultimo aggiornamento 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 ```"]]
