Eseguire 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 descritti in questo documento, devi creare un Google Cloud account, installare Google Cloud CLI e configurare il comando gcloud. Per saperne di più, consulta Configurare l'ambiente Cloud TPU.

Ruoli obbligatori

Per ottenere le autorizzazioni necessarie per creare una TPU e connetterti tramite 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 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: 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 Aree geografiche 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 acceleratori 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 \
     --project=$PROJECT_ID \
     --zone=$ZONE \
     --accelerator-type=$ACCELERATOR_TYPE \
     --version=$RUNTIME_VERSION

Connettersi alla VM Cloud TPU

Connettiti alla VM TPU tramite SSH utilizzando il comando seguente:

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 Connettersi a una VM TPU senza un indirizzo IP pubblico.

Installare 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

Verificare che PyTorch possa accedere alle TPU

Utilizza il comando seguente per verificare che PyTorch possa accedere alle 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 corrente 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')

Libera spazio

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

  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 la 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 comando seguente. Assicurati che la TPU non sia più elencata. L'eliminazione potrebbe richiedere alcuni minuti.

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

Passaggi successivi

Scopri di più sulle VM Cloud TPU: