Esegui 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 che mostra come addestrare un modello su una Cloud TPU, consulta uno dei tutorial di PyTorch per Cloud TPU.

Prima di iniziare

Prima di seguire questa guida rapida, devi creare un account per Google Cloud Platform, installare Google Cloud CLI e configurare il comando gcloud. Per ulteriori informazioni, vedi Configurare un account e un progetto Cloud TPU.

Crea una Cloud TPU con gcloud

Per creare una VM TPU nel progetto utente predefinito, nella rete e nella zona, esegui:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-8 \
   --version=tpu-ubuntu2204-base

Descrizioni flag comando

zone
La zona in cui prevedi di creare la Cloud TPU.
accelerator-type
Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.
version
La versione del software Cloud TPU.

Durante la creazione della TPU, puoi passare i flag --network e --subnetwork aggiuntivi se vuoi specificare la rete e la subnet predefinite. Se non vuoi utilizzare la rete predefinita, devi passare il flag --network. Il flag --subnetwork è facoltativo e può essere utilizzato per specificare una subnet predefinita per qualsiasi rete stai utilizzando (predefinita o specificata dall'utente). Per i dettagli su questi flag, consulta la gcloud pagina di riferimento delle API.

Connettiti alla VM Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central2-b

Installa PyTorch/XLA sulla tua VM TPU

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

Imposta la configurazione del 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, quindi copia e incolla il seguente script al suo interno.

    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 relativi alle risorse utilizzate in questa pagina, segui questi passaggi:

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

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

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

Passaggi successivi

Scopri di più sulle VM Cloud TPU: