Guida dell'utente di Cloud TPU PyTorch/XLA

Importante: puoi utilizzare le TPU utilizzando due architetture diverse: nodi TPU e VM TPU. Questo tutorial presuppone che tu stia utilizzando VM TPU. Le due architetture VM sono descritte in Architettura del sistema. A meno che tu non abbia bisogno di utilizzare nodi TPU, ti consigliamo di utilizzare le VM TPU. Per ulteriori informazioni sull'utilizzo dei PyTorch sui pod TPU mediante i nodi TPU, consulta l'articolo Addestramento dei modelli PyTorch sui pod Cloud TPU.

Esegui carichi di lavoro ML con PyTorch/XLA

Questa guida ti guida attraverso l'esecuzione di un semplice calcolo su una TPU v4 utilizzando PyTorch.

Configurazione di base

1. Crea una VM TPU con una TPU v4 che esegue il runtime della VM TPU per Pytorch 2.0:

     gcloud compute tpus tpu-vm create your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-v4-pt-2.0
   

   Nota: se vuoi installare una versione diversa di PyTorch, sostituisci tpu-vm-v4-pt-2.0 con la versione del runtime TPU che vuoi utilizzare. Per un elenco dei runtime di TPU con PyTorch, consulta la pagina sulle versioni supportate da PyTorch.

2. Connettiti alla VM TPU tramite SSH:

     gcloud compute tpus tpu-vm ssh your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8
   

   Punto chiave: durante questa guida, un prefisso di (vm) $ significa che devi eseguire il comando sull'istanza VM TPU.

3. Imposta la configurazione del dispositivo PJRT o XRT TPU.

   Nota: esistono due opzioni di runtime PyTorch/XLA: PJRT e XRT. Ti consigliamo di utilizzare PJRT, a meno che tu non abbia un motivo per utilizzare XRT. Per saperne di più sulle diverse configurazioni di runtime e vedere che hai un motivo per utilizzare XRT, consulta la documentazione di runtime PJRT.

   PJRT

       (vm)$ export PJRT_DEVICE=TPU
    

   XRT

       (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"
    

4. Se esegui l'addestramento con Cloud TPU v4, imposta anche la seguente variabile di ambiente:

     (vm)$ export TPU_NUM_DEVICES=4
   

Eseguire un semplice calcolo

1. Avvia l'interprete Python sulla VM TPU:

   (vm)$ python3
   

2. Importa i seguenti pacchetti PyTorch:

   import torch
   import torch_xla.core.xla_model as xm
   

3. Inserisci il seguente script:

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

   Nota: il messaggio di errore OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op è previsto e verrà corretto in una release futura.

   Viene visualizzato il seguente output:

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

Esecuzione di Resnet su una TPU di un singolo dispositivo

A questo punto, puoi eseguire il codice PyTorch / XL che preferisci. Ad esempio, puoi eseguire un modello ResNet con dati falsi:

(vm)$ git clone --recursive https://github.com/pytorch/xla.git
(vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

L'esempio ResNet viene addestrato per 1 periodo e richiede circa 7 minuti. Restituisce un output simile al seguente:

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

Al termine dell'addestramento ResNet, elimina la VM TPU.

(vm)$ exit

$ gcloud compute tpus tpu-vm delete tpu-name \
--zone=zone

L'eliminazione può richiedere qualche minuto. Verifica che le risorse siano state eliminate eseguendo gcloud compute tpus list --zone=${ZONE}.

Argomenti avanzati

Per i modelli con allocazioni frequenti e di grandi dimensioni, tcmalloc migliora le prestazioni rispetto alla funzione di runtime C/C++ malloc. Il valore predefinito di malloc utilizzato sulla VM TPU è tcmalloc. Puoi forzare il software VM VM a utilizzare lo standard malloc annullando l'impostazione della variabile di ambiente LD_PRELOAD:

   (vm)$ unset LD_PRELOAD

Negli esempi precedenti (il semplice calcolo e ResNet50), il programma PyTorch/XLA avvia il server XRT locale nello stesso processo dell'interprete Python. Puoi anche scegliere di avviare il servizio locale XRT in un processo separato:

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

Nota: la porta specificata in precedenza è la stessa specificata nel comando export XRT_TPU_CONFIG="localservice;0;localhost:51011".

Il vantaggio di questo approccio è che la cache delle compilation persiste durante le esecuzioni di addestramento. Quando esegui il server XLA in un processo separato, le informazioni di logging lato server vengono scritte in /tmp/xrt_server_log.

(vm)$ ls /tmp/xrt_server_log/

server_20210401-031010.log

Profilazione delle prestazioni delle VM TPU

Per ulteriori informazioni sulla profilazione dei modelli su VM TPU, consulta la pagina Profilazione delle prestazioni PYTorch XLA.

Esempi di pod PyTorch/XLA TPU

Vedi Pod VM PyTorch TPU per informazioni sulla configurazione ed esempi per l'esecuzione di PyTorch/XLA su un pod VM TPU.

Docker su VM VM

Questa sezione mostra come eseguire Docker su VM TPU con PyTorch/XLA preinstallato.

Nota: al momento, l'esecuzione di Docker su VM VM è supportata solo su TPU per singolo dispositivo (v2-8, v3-8).

Immagini Docker disponibili

Puoi fare riferimento a GitHub README per trovare tutte le immagini Docker VM disponibili.

Esegui immagini Docker su VM VM

Nota: l'ultimo output comando, root:/#, è quello che viene visualizzato dopo aver eseguito l'immagine Docker.

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

Nota: devi eseguire il container utilizzando il flag --privileged per assicurarti che i dispositivi TPU siano visibili all'interno del container. In questo container è preinstallato un ambiente conda pytorch. Ogni comando che inizia con (pytorch) deve essere eseguito all'interno del container.

Verifica libtpu

Nota:libtpu dovrebbe essere preinstallato sulle immagini Docker di VMVM più recenti, ma è sempre meglio verificarlo manualmente.

Per verificare che libtpu sia installato, esegui:

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

L'output dovrebbe essere simile al seguente:

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

Se non viene visualizzato alcun risultato, puoi installare manualmente il file libtpu corrispondente utilizzando:

(pytorch) root:/# pip install torch_xla[tpuvm]

Verifica tcmalloc

tcmalloc è il Malloc predefinito che utilizziamo sulla VM TPU. Per ulteriori informazioni, leggi questa sezione. Questa libreria dovrebbe essere preinstallata sulle immagini Docker VM più recenti, ma è sempre meglio verificarla manualmente. Puoi eseguire questo comando per verificare che la libreria sia installata.

(pytorch) root:/# echo $LD_PRELOAD

Questo dovrebbe generare un output simile al seguente:

/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4

Se LD_PRELOAD non è impostato, puoi eseguire manualmente:

(pytorch) root:/# sudo apt-get install -y google-perftools
(pytorch) root:/# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4"

Verifica dispositivo

Per verificare che la VM TPU sia disponibile, esegui questo comando:

(pytorch) root:/# ls /dev | grep accel

Dovrebbero essere generati i seguenti risultati.

accel0
accel1
accel2
accel3

Se non vengono mostrati risultati, molto probabilmente non hai avviato il container con il flag --privileged.

Esegui un modello

Per verificare se il dispositivo VM TPU è disponibile, esegui questo comando:

(pytorch) root:/# export XRT_TPU_CONFIG="localservice;0;localhost:51011"
(pytorch) root:/# python3 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --num_epochs 1