Esecuzione del codice PyTorch nelle slice 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 fare lo scale up del codice eseguendolo su una slice TPU. Le slice TPU sono più schede TPU collegate tra loro tramite connessioni di rete dedicate ad alta velocità. Questo documento è un'introduzione all'esecuzione del codice PyTorch sulle slice TPU.

Ruoli obbligatori

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

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.

Creare una slice Cloud TPU

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

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

    Descrizioni delle variabili di ambiente

    • PROJECT_ID: l'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 saperne di più sui tipi di acceleratore supportati per ogni versione TPU, consulta Versioni TPU.
    • RUNTIME_VERSION: la versione software di Cloud TPU.

  2. Crea la VM TPU eseguendo il comando seguente:

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

Installare PyTorch/XLA sulla slice

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

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.

  1. Installa PyTorch/XLA su tutti i worker delle 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 delle 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"

Eseguire uno script di addestramento sulla 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 saperne di più 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%

Libera spazio

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:

    (vm)$ 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  \
    --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 comando seguente non deve includere nessuna delle risorse create in questo tutorial:

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