Eseguire un calcolo su una VM Cloud TPU utilizzando PyTorch

Questa guida rapida mostra come creare una Cloud TPU, installare PyTorch ed eseguire un semplice calcolo su una Cloud TPU. Per un tutorial più approfondito come addestrare un modello su una Cloud TPU, consulta uno dei Tutorial su Cloud TPU PyTorch.

Prima di iniziare

Prima di seguire questa guida rapida, devi creare un account Google Cloud, installare Google Cloud CLI e configurare il comando gcloud. Per saperne di più, consulta la sezione Configurare un account e un progetto Cloud TPU.

Crea una Cloud TPU con gcloud

.

Per creare una VM TPU nel progetto, nella rete e nella zona utente predefiniti, eseguire:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Descrizioni dei flag dei comandi

zone
La zona in cui prevedi di creare la tua Cloud 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.
version
La versione del software Cloud TPU.

Durante la creazione della TPU, puoi trasmettere --network e --subnetwork flag se devi specificare la rete e la subnet predefinite. Se non vuoi utilizzare la rete predefinita, devi passare il --network flag. Il flag --subnetwork è facoltativo e può essere utilizzato per specificare una subnet predefinita per qualsiasi rete in uso (predefinita o specificata dall'utente). Consulta gcloud pagina di riferimento delle API per informazioni dettagliate su questi flag.

Connettiti alla VM Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

Installare PyTorch/XLA sulla VM TPU

   (vm)$ pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

Imposta configurazione di runtime TPU

Assicurati che il runtime PyTorch/XLA utilizzi la TPU.

   (vm) $ export PJRT_DEVICE=TPU

Esegui un semplice calcolo:

  1. Crea un file denominato tpu-test.py nella directory corrente e copia e incolla al suo interno 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:

      (vm)$ 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 per le risorse utilizzate in questa pagina, segui questi passaggi.

  1. Disconnettiti dall'istanza Compute Engine, se non lo hai già fatto Fatto:

    (vm)$ 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 \
  --zone=us-central1-a

L'output di questo comando dovrebbe confermare che la TPU è stata eliminata.

Passaggi successivi

Scopri di più sulle VM Cloud TPU: