Eseguire il codice PyTorch sulle sezioni del pod di TPU

Configurare un pod di una VM TPU che esegue PyTorch ed eseguire un calcolo

PyTorch/XLA richiede che tutte le VM TPU siano in grado di accedere al codice e ai dati del modello. Uno script di avvio scarica il software necessario per distribuire i dati del modello a tutte le VM TPU.

In Cloud Shell, esegui il comando seguente per assicurarti di eseguire la versione più recente di gcloud:
```
$ gcloud components update
```
Se devi installare gcloud, utilizza il comando seguente:
```
$ sudo apt install -y google-cloud-sdk
```
Esporta le seguenti variabili di ambiente:
```
$ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=zone
$ export RUNTIME_VERSION=tpu-vm-pt-1.13
```
Nota: una versione predefinita del runtime Cloud TPU è indicata come RUNTIME_VERSION. Puoi eseguire qualsiasi versione PyTorch supportata.
Crea la VM TPU
```
$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone ${ZONE} --project ${PROJECT_ID} --accelerator-type v3-32 \
--version ${RUNTIME_VERSION} --metadata startup-script='#! /bin/bash
cd /usr/share/
git clone --recursive https://github.com/pytorch/pytorch
cd pytorch/
git clone --recursive https://github.com/pytorch/xla.git
EOF'
```
Lo script di avvio genera la directory /usr/share/pytorch/xla sulla VM TPU e scarica il codice del modello in questa directory. Il download è rapido, tuttavia l'elaborazione dei file può richiedere diversi minuti. Puoi verificare che l'elaborazione del file sia stata completata eseguendo ls -al in /usr/share/pytorch per vedere se è stata creata la directory xla.
Al termine dell'elaborazione del file, esegui il comando ssh seguente per generare le chiavi ssh su ssh tra i worker VM sul pod. Quindi, avvia il corso di formazione.

Nota: quando ssh in un pod di TPU, per impostazione predefinita ti connetterai alla prima VM TPU (worker 0).

Per ssh in altre VM TPU associate al pod TPU, aggiungi --worker ${WORKER_NUMBER} nel comando ssh, dove WORKER_NUMBER è un indice basato su 0. Per ulteriori dettagli, consulta la guida dell'utente sulle VM TPU.
```
$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
 --zone ${ZONE} \
 --project ${PROJECT_ID}
```
Punto chiave: da questo punto in poi, un prefisso (vm) $ indica che devi eseguire il comando sull'istanza VM di Compute Engine.

Aggiorna i metadati ssh del progetto:
```
(vm)$ gcloud compute config-ssh
```
```
(vm)$ export TPU_NAME=tpu-name
```
```
(vm)$ python3 -m torch_xla.distributed.xla_dist \
 --tpu=${TPU_NAME} -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py \
 --fake_data --model=resnet50 --num_epochs=1
```
L'addestramento richiede circa 3 minuti. Al termine, dovresti visualizzare un messaggio simile al seguente:
```
Epoch 1 test end 23:49:15, Accuracy=100.00
10.164.0.11 [0] Max Accuracy: 100.00%
```
Puoi trovare il logging lato server in /tmp/xrt_server_log su ciascun worker.
```
(vm)$ ls /tmp/xrt_server_log/
```
```
server_20210401-031010.log
```
Se vuoi riavviare XRT_SERVER (nel caso in cui il server si trovi in uno stato non integro), puoi superare --restart-tpuvm-pod-server quando esegui xla_dist. Le nuove impostazioni del server XRT, come le variabili di ambiente, come LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4, hanno effetto solo dopo aver riavviato il server.

Esegui la pulizia

Al termine della VM TPU, segui questi passaggi per eseguire la pulizia delle risorse.

Disconnettiti da Compute Engine:
```
(vm)$ exit
```

Elimina Cloud TPU.

$ gcloud compute tpus tpu-vm delete tpu-name \
  --zone europe-west4-a

Verifica che le risorse siano state eliminate eseguendo il comando seguente. Assicurati che la tua TPU non sia più presente nell'elenco. L'eliminazione può richiedere qualche minuto.
```
$ gcloud compute tpus tpu-vm list \
  --zone europe-west4-a
```

Pod con coordinatore remoto

Ti consigliamo di utilizzare un coordinatore remoto per consentire il ripristino automatico del pod quando si verifica un evento di manutenzione di TPU. Un coordinatore remoto è una VM standard di Compute Engine, non una VM TPU. Emette comandi per le VM TPU nel tuo pod.

Esporta le seguenti variabili di ambiente:

$ export TPU_NAME=tpu-name
$ export ZONE=zone
$ export PROJECT_ID=project-id

Crea una sezione del pod VM TPU:

$ gcloud compute tpus tpu-vm create tpu-name \
 --zone europe-west4-a --project project-id --accelerator-type v3-32 \
 --version tpu-vm-pt-1.13 --metadata startup-script='#! /bin/bash
cd /usr/share/
git clone --recursive https://github.com/pytorch/pytorch
cd pytorch/
git clone --recursive https://github.com/pytorch/xla.git
EOF'

Esporta le seguenti variabili di ambiente necessarie per creare la VM Coordinator remoto (RC):
```
$ export RC_NAME=rc-name
$ export ZONE=zone
```
Crea una VM del coordinatore remoto eseguendo:
```
(vm)$ gcloud compute --project=project-id instances create rc-name \
 --zone=zone  \
 --machine-type=n1-standard-1  \
 --image-family=torch-xla \
 --image-project=ml-images  \
 --boot-disk-size=200GB \
 --scopes=https://www.googleapis.com/auth/cloud-platform
```
Al termine dell'esecuzione del comando gcloud compute, verifica che il prompt di Cloud Shell sia cambiato da username@projectname a username@vm-name. Questa modifica mostra che hai eseguito l'accesso alla VM del coordinatore remoto.

ssh nell'istanza del coordinatore remoto:

(vm)$ gcloud compute ssh rc-name --zone=zone

Attiva l'ambiente torch-xla-1.13 ed esegui l'addestramento da lì.

(vm)$ gcloud compute config-ssh

(vm)$ conda activate torch-xla-1.13
(vm)$ python3 -m torch_xla.distributed.xla_dist \
 --tpu=tpu-name \
 --restart-tpuvm-pod \
 --env LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py \
 --fake_data \
 --model=resnet50 \
 --num_epochs=1

L'addestramento richiede circa 3 minuti e genera un messaggio simile al seguente:

Epoch 1 test end 23:19:53, Accuracy=100.00
Max Accuracy: 100.00%

Al termine dell'addestramento ResNet, esci dalla VM TPU ed elimina la VM del coordinatore remoto e la VM TPU.

(vm)$ exit

$ gcloud compute instances delete rc-name  \
  --zone=zone

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