Esegui un calcolo su una VM Cloud TPU utilizzando PyTorch

Questo documento fornisce una breve introduzione all'utilizzo di PyTorch e Cloud TPU.

Prima di iniziare

Prima di eseguire i comandi in questo documento, devi creare un account Google Cloud , installare Google Cloud CLI e configurare il comando gcloud. Per maggiori informazioni, consulta Configura l'ambiente Cloud 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 nel 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.

Crea una Cloud TPU utilizzando gcloud

  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=us-east5-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

    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 ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
    • RUNTIME_VERSION: la versione software di Cloud TPU.

  2. Crea la tua VM TPU eseguendo questo comando:

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

Connettiti alla VM Cloud TPU

Connettiti alla VM TPU tramite SSH utilizzando il seguente comando:

gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Se non riesci a connetterti a una VM TPU tramite SSH, è possibile che la VM TPU non abbia un indirizzo IP esterno. Per accedere a una VM TPU senza un indirizzo IP esterno, segui le istruzioni riportate in Connettiti a una VM TPU senza un indirizzo IP pubblico.

Installa PyTorch/XLA sulla VM TPU

  1. Salva quanto segue in un file denominato requirements.txt:

    --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
ray[default]==2.40.0

  2. Installa le dipendenze:

    pip install -r requirements.txt

Verifica che PyTorch possa accedere alle TPU

Utilizza il seguente comando per verificare che PyTorch possa accedere alle tue TPU:

PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

L'output del comando dovrebbe essere simile al seguente:

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Eseguire un calcolo di base

  1. Crea un file denominato tpu-test.py nella directory attuale e aggiungi il seguente script:

    import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

  2. Esegui lo script:

    PJRT_DEVICE=TPU python3 tpu-test.py

    L'output dello script mostra il risultato del calcolo:

    tensor([[-0.2121,  1.5589, -0.6951],
         [-0.7886, -0.2022,  0.9242],
         [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questa pagina, segui questi passaggi.

  1. Se non l'hai ancora fatto, disconnettiti dall'istanza Cloud TPU:

    exit

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

  2. Elimina la tua Cloud TPU.

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

  3. Verifica che le risorse siano state eliminate eseguendo il seguente comando. Assicurati che la TPU non sia più elencata. L'eliminazione potrebbe richiedere diversi minuti.

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

Passaggi successivi

Scopri di più sulle VM Cloud TPU: