Esecuzione del codice PyTorch sulle sezioni TPU

Attenzione:questa pagina si applica all'API Cloud TPU, che non è più in fase di sviluppo attivo. L'API Cloud TPU riceverà solo correzioni di bug e aggiornamenti della sicurezza. Per le funzionalità più recenti e il supporto delle versioni TPU più recenti, ti consigliamo di utilizzare Compute Engine o Google Kubernetes Engine per gestire le risorse TPU. Per saperne di più, consulta Risorse TPU in Compute Engine e Informazioni sulle TPU in GKE. L'API Cloud TPU non è supportata su TPU7x e versioni TPU successive.

Prima di eseguire i comandi in questo documento, segui le istruzioni riportate in Configurare un account e un progetto Cloud TPU.

Dopo aver eseguito il codice PyTorch su una singola VM TPU, puoi fare lo scale up il codice eseguendolo su una sezione TPU. Le sezioni TPU sono più schede TPU connesse tra loro tramite connessioni di rete dedicate ad alta velocità. Questo documento introduce l'esecuzione del codice PyTorch sulle sezioni TPU.

Ruoli obbligatori

Per ottenere le autorizzazioni necessarie per creare una TPU e connetterti a quest'ultima utilizzando SSH, chiedi all'amministratore di concederti i seguenti ruoli IAM sul progetto:

• Amministratore TPU (roles/tpu.admin)
• Utente Service Account (roles/iam.serviceAccountUser)
• Compute Viewer (roles/compute.viewer)

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Crea uno slice Cloud TPU

1. Definisci le variabili di ambiente per i parametri dei comandi Google Cloud CLI.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-central1-b
   export ACCELERATOR_TYPE=v6e-8
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   Descrizioni delle variabili di ambiente

   • PROJECT_ID: il tuo ID progetto Google Cloud . 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 maggiori informazioni sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
   • RUNTIME_VERSION: la versione software di Cloud TPU.

2. Crea la VM TPU eseguendo il seguente comando:

   gcloud compute tpus tpu-vm create $TPU_NAME \
        --zone=$ZONE \
        --project=$PROJECT_ID \
        --accelerator-type=$ACCELERATOR_TYPE \
        --version=$RUNTIME_VERSION
   

Installa PyTorch/XLA sulla tua sezione

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 --command.

1. Crea un file denominato requirements.txt con i seguenti contenuti:

   --find-links https://storage.googleapis.com/libtpu-releases/index.html
   --find-links https://storage.googleapis.com/libtpu-wheels/index.html
   torch~=2.6.0
   torch_xla[tpu]~=2.6.0
   torchvision
   ray[default]==2.40.0
   

2. Copia requirements.txt in ogni VM della tua sezione:

   gcloud compute tpus tpu-vm tpu-vm scp ./requirements.txt \
        $TPU_NAME:~/ \
        --zone=$ZONE \
        --project=$PROJECT_ID \
        --worker=all
   

3. Installa le dipendenze su ogni VM nella tua slice:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
        --zone=$ZONE \
        --project=$PROJECT_ID \
        --worker=all \
        --command="pip3 install -r requirements.txt"
   

4. Clona XLA su tutti i worker VM 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"
   

Esegui uno script di addestramento sulla tua slice TPU

Esegui lo script di addestramento su tutti i worker. Lo script di addestramento utilizza una strategia di sharding Single Program Multiple Data (SPMD). Per ulteriori informazioni su SPMD, consulta la guida dell'utente di PyTorch/XLA SPMD.

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"

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 liberare spazio dalle risorse.

1. Disconnettiti dall'istanza Cloud TPU, se non l'hai già fatto:

   exit
   

   Il tuo prompt dovrebbe ora essere username@projectname, a indicare che ti trovi in Cloud Shell.

2. Elimina le risorse Cloud TPU.

   gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE
   

3. 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:

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