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

Prima di iniziare

Prima di seguire questa guida rapida, devi creare un account Google Cloud Platform, installare Google Cloud CLI e configurare il comando gcloud. Per saperne di più, consulta 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 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 --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). Per informazioni dettagliate su questi flag, consulta la gcloud pagina di riferimento dell'API.

Connettiti alla VM Cloud TPU

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

Installa PyTorch/XLA sulla tua VM TPU

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

Impostare la configurazione di runtime della TPU

Assicurati che il runtime PyTorch/XLA utilizzi la TPU.

   (vm) $ export PJRT_DEVICE=TPU

Verificare che PyTorch possa accedere alle TPU

  1. Crea un file denominato tpu-count.py nella directory corrente e copia e incolla al suo interno il seguente script.

    import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')

  2. Esegui lo script:

    (vm)$ python3 tpu-count.py

    L'output dello script mostra il risultato del calcolo:

    PyTorch can access 8 TPU cores

Eseguire un calcolo di base

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

  1. Se non l'hai ancora fatto, disconnettiti dall'istanza Compute Engine:

    (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: