Esegui il codice PyTorch nelle sezioni del pod TPU

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

Dopo aver eseguito il codice PyTorch su una singola VM TPU, puoi eseguirne l'upgrade su una sezione di pod TPU. Le sezioni del pod 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 nelle sezioni di pod TPU.

Crea una sezione del pod Cloud TPU

  1. Definisci alcune variabili di ambiente per semplificare l'utilizzo dei comandi.

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

    Descrizioni delle variabili di ambiente

    PROJECT_ID
    Il tuo Google Cloud ID progetto.
    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 di TPU.
    ZONE
    La zona in cui prevedi di creare la tua Cloud TPU.
    RUNTIME_VERSION
    La versione del software Cloud TPU.
    TPU_NAME
    Il nome assegnato dall'utente alla tua 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 nel tuo slice del pod

Dopo aver creato la sezione del pod TPU, devi installare PyTorch su tutti gli host della sezione del pod TPU. Puoi farlo utilizzando il comando gcloud compute tpus tpu-vm ssh con i parametri --worker=all e --commamnd.

  1. Installa PyTorch/XLA su tutti i worker VM 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. 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 nella sezione del pod 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.

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

  1. Se non l'hai già fatto, scollega l'istanza Compute Engine:

    (vm)$ exit

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

  2. Elimina le risorse Cloud TPU e Compute Engine.

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

  3. Verifica che le risorse siano state eliminate eseguendo gcloud compute tpus execution-groups 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}