Exécuter le code PyTorch sur des tranches de pod TPU

Configurer un pod de VM TPU exécutant PyTorch et effectuer un calcul

PyTorch/XLA nécessite que toutes les VM TPU puissent accéder au code et aux données du modèle. Un script de démarrage télécharge le logiciel nécessaire pour distribuer les données du modèle à toutes les VM TPU.

Dans Cloud Shell, exécutez la commande suivante pour vous assurer que vous exécutez la version la plus récente de gcloud :
```
$ gcloud components update
```
Si vous devez installer gcloud, utilisez la commande suivante:
```
$ sudo apt install -y google-cloud-sdk
```
Exportez les variables d'environnement suivantes :
```
$ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=zone
$ export RUNTIME_VERSION=tpu-vm-pt-1.13
```
Remarque:Une version d'exécution par défaut de Cloud TPU est indiquée par la valeur RUNTIME_VERSION. Vous pouvez exécuter n'importe quelle version PyTorch compatible.
Créez 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'
```
Le script de démarrage génère le répertoire /usr/share/pytorch/xla sur la VM TPU et télécharge le code du modèle dans ce répertoire. Le téléchargement est rapide, mais le traitement du fichier peut prendre plusieurs minutes. Vous pouvez vérifier que le traitement du fichier est terminé en exécutant ls -al dans /usr/share/pytorch pour voir si le répertoire xla a été créé.
Une fois le traitement du fichier terminé, exécutez la commande ssh suivante pour générer des clés SSH vers ssh entre les nœuds de calcul de VM du pod. Vous pouvez ensuite commencer l'entraînement.

Remarque : Lorsque vous ssh dans un pod TPU, vous vous connectez par défaut à la première VM TPU (nœud de calcul 0).

Pour ssh dans les autres VM TPU associées au pod TPU, ajoutez --worker ${WORKER_NUMBER} dans la commande ssh, où WORKER_NUMBER est un index basé sur 0. Pour en savoir plus, consultez le Guide utilisateur de la VM TPU.
```
$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
 --zone ${ZONE} \
 --project ${PROJECT_ID}
```
Point clé : À partir de ce point, le préfixe (vm) $ signifie que vous devez exécuter la commande sur l'instance de VM Compute Engine.

Mettez à jour les métadonnées ssh du projet :
```
(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'entraînement dure environ trois minutes. Une fois l'opération terminée, un message semblable au suivant doit s'afficher :
```
Epoch 1 test end 23:49:15, Accuracy=100.00
10.164.0.11 [0] Max Accuracy: 100.00%
```
Vous pouvez trouver la journalisation côté serveur dans /tmp/xrt_server_log sur chaque nœud de calcul.
```
(vm)$ ls /tmp/xrt_server_log/
```
```
server_20210401-031010.log
```
Si vous souhaitez redémarrer XRT_SERVER (si le serveur n'est pas opérationnel), vous pouvez transmettre --restart-tpuvm-pod-server lors de l'exécution de xla_dist. Les nouveaux paramètres de serveur XRT tels que les variables d'environnement telles que LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 ne prendront effet qu'au redémarrage du serveur.

Effectuer un nettoyage

Lorsque vous avez fini d'utiliser votre VM TPU, procédez comme suit pour nettoyer vos ressources.

Déconnectez-vous de Compute Engine:
```
(vm)$ exit
```

Supprimez votre Cloud TPU.

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

Vérifiez que les ressources ont bien été supprimées en exécutant la commande suivante. Assurez-vous que votre TPU n'est plus répertorié. La suppression peut prendre plusieurs minutes.
```
$ gcloud compute tpus tpu-vm list \
  --zone europe-west4-a
```

Pods avec coordinateur à distance

Nous vous recommandons d'utiliser un coordinateur à distance pour permettre la récupération automatique de votre pod lorsqu'un événement de maintenance TPU se produit. Un coordinateur distant est une VM Compute Engine standard, et non une VM TPU. Il envoie des commandes aux VM TPU de votre pod.

Exportez les variables d'environnement suivantes :

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

Créez une tranche de pod de 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'

Exportez les variables d'environnement suivantes pour créer la VM du coordinateur distant:
```
$ export RC_NAME=rc-name
$ export ZONE=zone
```
Créez une VM de coordinateur distant en exécutant la commande suivante :
```
(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
```
Une fois l'exécution de la commande gcloud compute terminée, vérifiez que l'invite de l'interface système est passée de username@projectname à username@vm-name. Cette modification indique que vous êtes maintenant connecté à la VM du coordinateur distant.

ssh dans l'instance du coordinateur distant :

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

Activez l'environnement torch-xla-1.13, puis exécutez votre entraînement.

(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'entraînement dure environ trois minutes et génère un message semblable à celui-ci :

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

Une fois l'entraînement ResNet terminé, quittez la VM TPU, et supprimez la VM du coordinateur à distance et la VM TPU.

(vm)$ exit

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

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