Guide de l'utilisateur de Cloud TPU PyTorch/XLA

Exécuter des charges de travail de ML avec PyTorch/XLA

Avant d'entamer les procédures de ce guide, configurez une VM TPU et une ssh en procédant comme décrit dans la section Préparer un projet GCP. Les ressources nécessaires à l'exécution des commandes de ce guide seront alors configurées.

PyTorch 1.8.1 et PyTorch/XLA 1.8.1 sont préinstallés sur la VM TPU.

Configuration de base

Définissez la configuration de l'appareil TPU XRT:

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

Pour les modèles comportant des allocations fréquentes et importantes, tcmalloc améliore considérablement les performances par rapport à la fonction d'exécution C/C++ malloc. La valeur par défaut de malloc utilisée sur la VM TPU est tcmalloc. Vous pouvez forcer l'environnement d'exécution de la VM TPU à utiliser la version standard malloc en désactivant la variable d'environnement LD_PRELOAD:

   (vm)$ unset LD_PRELOAD

Modifier la version de PyTorch

Si vous ne souhaitez pas utiliser la version PyTorch préinstallée sur les VM TPU, installez-la. Exemple :

(vm)$ sudo bash /var/scripts/docker-login.sh
(vm)$ sudo docker rm libtpu || true
(vm)$ sudo docker create --name libtpu gcr.io/cloud-tpu-v2-images/libtpu:pytorch-1.9 "/bin/bash"
(vm)$ sudo docker cp libtpu:libtpu.so /lib
(vm)$ sudo pip3 uninstall --yes torch torch_xla torchvision
(vm)$ sudo pip3 install torch==1.9.0
(vm)$ sudo pip3 install torchvision==0.10.0
(vm)$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.9-cp38-cp38-linux_x86_64.whl

Effectuez un calcul simple

Démarrez l'interpréteur Python sur la VM TPU:
```
(vm)$ python3
```

Importez les packages PyTorch suivants:

import torch
import torch_xla.core.xla_model as xm

Saisissez le script suivant:

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

Le résultat suivant s'affiche:

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

Exécuter Resnet sur un TPU à appareil unique

À ce stade, vous pouvez exécuter n'importe quel code PyTorch / XLA. Par exemple, vous pouvez exécuter un modèle ResNet avec des données fictives:

(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'exemple ResNet est entraîné pendant 1 époque et prend environ 7 minutes. Elle renvoie un résultat semblable à celui-ci:

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

Une fois l'entraînement ResNet terminé, supprimez la VM TPU.

(vm)$ exit

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

La suppression peut prendre plusieurs minutes. Vérifiez que les ressources ont été supprimées en exécutant gcloud alpha compute tpus list --zone=${ZONE}.

Configuration avancée

Dans les exemples ci-dessus (le calcul simple et ResNet50), votre programme PyTorch/XLA démarre le serveur XRT local selon le même processus que l'interpréteur Python. Vous pouvez également choisir de démarrer le service local XRT dans un processus distinct:

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

L'avantage de cette approche est que le cache de compilation conserve les sessions d'entraînement. Lorsque vous exécutez le serveur XLA dans un processus distinct, les informations de journalisation côté serveur sont écrites dans /tmp/xrt_server_log.

(vm)$ ls /tmp/xrt_server_log/

server_20210401-031010.log

Profilage des performances de VM TPU

Pour en savoir plus sur le profilage de vos modèles sur une VM TPU, consultez la page Profilage des performances PyTorch XLA.

Pods

PyTorch/XLA exige que toutes les VM TPU puissent accéder au code et aux données du modèle. Pour ce faire, vous pouvez utiliser le script de démarrage suivant lorsque vous créez le pod de la VM TPU. Il effectue le téléchargement des données sur 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:
```
$ sudo /opt/google-cloud-sdk/bin/gcloud components update
```

Exportez les variables d'environnement suivantes :

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

Créer la VM TPU

$ gcloud alpha 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'

À mesure que vous suivez ces instructions, exécutez chaque commande commençant par (vm)$ dans l'interface système de votre VM TPU.

ssh à n'importe quel travailleur TPU (par exemple, le nœud de calcul 0). Vérifiez si le téléchargement des données ou du modèle est terminé. Le téléchargement est généralement très 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 le répertoire de téléchargement. Les fichiers sont téléchargés dans l'ordre que vous avez spécifié dans le script de démarrage. Exemple :
```
download file1
download file2
download file3
```
Si file3 est le dernier fichier du script de démarrage et qu'il est affiché par la commande ls, le traitement du fichier est terminé.

Une fois le traitement du fichier terminé, générez des clés SSH vers ssh entre les nœuds de calcul de VM sur le pod. Vous pouvez ensuite commencer la formation.

Remarque :Pour les pods TPU, la valeur par défaut est ssh dans la première VM TPU. Pour ajouter ssh à d'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 de l'utilisateur de VM TPU.
```
(vm)$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
 --zone ${ZONE} \
 --project ${PROJECT_ID}
```
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. Un message semblable au suivant s'affiche alors:
```
Epoch 1 test end 23:49:15, Accuracy=100.00
10.164.0.11 [0] Max Accuracy: 100.00%
```
Vous trouverez la journalisation côté serveur sous /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. Notez que 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.

Une fois l'entraînement ResNet terminé, supprimez la VM TPU.

(vm)$ exit

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

Vérifiez que les ressources ont été supprimées en exécutant la commande gcloud alpha compute tpus list --zone=${ZONE}. La suppression peut prendre plusieurs minutes.

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 alpha compute tpus tpu-vm create tpu-name \
 --zone europe-west4-a --project tpu-prod-env-one-vm --accelerator-type v3-32 \
 --version v2-alpha --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, nécessaires pour créer la VM du coordinateur distant:
```
$ export VM_NAME=vm-name
$ export ZONE=zone
```
Créez une VM Google Coordinator distante en exécutant la commande suivante:
```
(vm)$ gcloud compute --project=project-id instances create vm-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 vm-name --zone=zone

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

(vm)$ gcloud compute config-ssh

(vm)$ conda activate torch-xla-1.8.1
(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'exécution de 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 à coordinateur distant et la VM TPU.

(vm)$ exit

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

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