Guide de l'utilisateur de Cloud TPU PyTorch/XLA

Important : Vous pouvez utiliser des TPU à l'aide de deux architectures différentes: les nœuds TPU et les VM TPU. Ce tutoriel suppose que vous utilisez des VM TPU. Les deux architectures de VM sont décrites dans la section Architecture système. À moins que vous ne sachiez utiliser des nœuds TPU, nous vous recommandons d'utiliser des VM TPU. Pour en savoir plus sur l'utilisation de PyTorch sur les pods TPU à l'aide de nœuds TPU, consultez le guide des utilisateurs XLA pour PyTorch.

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

Avant de commencer les procédures de ce guide, configurez une VM TPU et un ssh comme décrit dans la section Préparer un projet Google Cloud.

Point clé : Dans ce guide, le préfixe (vm) $ signifie que vous devez exécuter la commande sur l'instance de VM TPU. L'invite $ signifie que vous devez exécuter la commande dans votre interface système locale.

Deux versions de logiciel logiciel sont disponibles pour Pytorch / XLA:

1. Dernière image disponible: tpu-vm-pt-1.11
2. tpu-vm-pt-2.10
3. v2 alpha

L'image tpu-vm-pt-1.11 est préinstallée avec PyTorch 1.11.0 et Pytorch / XLA 1.11.0. L'image tpu-vm-pt-1.10 est préinstallée avec PyTorch 1.10.0 et Pytorch / XLA 1.10.0. L'image v2-alpha est préinstallée avec PyTorch 1.8.1.

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 avec des allocations importantes et fréquentes, tcmalloc améliore 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 le logiciel de VM TPU à utiliser la classe malloc standard 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 version que vous souhaitez utiliser. Par exemple, si vous souhaitez utiliser 1.9 :

(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

Effectuer un calcul simple

1. Démarrez l'interpréteur Python sur la VM TPU :

   (vm)$ python3
   

2. Importez les packages PyTorch suivants :

   import torch
   import torch_xla.core.xla_model as xm
   

3. 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)
   

   Remarque : Le message d'erreur OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op est attendu et sera corrigé dans une version ultérieure.

   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 individuel

À 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 une époque et prend environ sept minutes. La commande renvoie un résultat semblable au suivant :

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 la commande gcloud alpha compute tpus list --zone=${ZONE}.

Configuration avancée

Dans les exemples précédents (le calcul simple et ResNet50), le 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

Remarque : Le port spécifié précédemment est le même que celui que vous avez spécifié précédemment dans la commande export XRT_TPU_CONFIG="localservice;0;localhost:51011".

L'avantage de cette approche est que le cache de compilation persiste pendant les exécutions d'entraînement. Lors de l'exécution du 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 des 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 nécessite que toutes les VM TPU puissent accéder au code et aux données du modèle. Utilisez le script de démarrage suivant lorsque vous créez le pod de VM TPU, qui télécharge les données sur toutes les VM TPU.

Remarque:Pour l'instant, l'exécution de docker sur une VM TPU n'est possible que sur des TPU individuels (v2-8, v3-8).

1. 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
   

2. 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.11
   

3. Créez 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'
   

   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.

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

4. 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 rapide, mais le traitement d'un 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 figure dans le résultat de 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 l'entraînement.

   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 la 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. 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 trouverez 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 du serveur XRT, tels que les variables d'environnement telles que LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4, ne prennent effet qu'au redémarrage du serveur.

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

   (vm)$ exit
   

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

6. 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.

1. Exportez les variables d'environnement suivantes :

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

2. 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 tpu-vm-pt-1.11 --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'
   

3. Exportez les variables d'environnement suivantes, nécessaires à la création de la VM du coordinateur distant :

   $ export VM_NAME=vm-name
   $ export ZONE=zone
   

4. Créez une VM de coordinateur distant 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.

5. ssh dans l'instance du coordinateur distant :

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

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

   (vm)$ gcloud compute config-ssh
   

   (vm)$ conda activate torch-xla-1.11
   (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%
   

7. 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 vm-name  \
     --zone=zone
   

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

Docker sur une VM TPU

Cette section explique comment exécuter Docker sur une VM TPU avec PyTorch/XLA préinstallé.

Remarque : Pour l'instant, l'exécution de docker sur une VM TPU n'est possible que sur des TPU individuels (v2-8, v3-8).

Images Docker disponibles

Vous pouvez vous reporter au README de GitHub pour trouver toutes les images Docker disponibles de la VM TPU.

Exécuter des images Docker sur une VM TPU

Remarque:Le dernier résultat de la commande, root:/#, est l'invite que vous voyez après l'exécution de l'image 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:/#

Remarque : Vous devez exécuter le conteneur à l'aide de l'option --privileged pour vous assurer que les appareils TPU sont visibles dans le conteneur. pytorch environnement de conda est préinstallé sur ce conteneur. Chaque commande commençant par (pytorch) doit être exécutée dans le conteneur.

Valider libtpu

Remarque:libtpu doit être préinstallé sur les Dockers de VM TPUVM plus récents, mais il est toujours préférable de le vérifier manuellement.

Pour vérifier que libtpu est installé, exécutez la commande suivante:

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

Vous devriez obtenir un résultat semblable à celui-ci:

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

Si aucun résultat ne s'affiche, vous pouvez installer manuellement libtpu à l'aide de la commande suivante:

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

Validez le domaine tcmalloc

tcmalloc est le Malloc par défaut que nous utilisons sur la VM TPU. Pour en savoir plus, consultez cette section. Cette bibliothèque doit être préinstallée sur les images Docker les plus récentes de la VM TPU, mais il est toujours préférable de la valider manuellement. Vous pouvez exécuter la commande suivante pour vérifier que la bibliothèque est installée.

(pytorch) root:/# echo $LD_PRELOAD

Vous devriez obtenir un résultat semblable à celui-ci:

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

Si LD_PRELOAD n'est pas défini, vous pouvez exécuter manuellement:

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

Valider l'appareil

Vous pouvez vérifier que l'appareil de VM TPU est disponible en exécutant la commande suivante:

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

Les résultats suivants devraient s'afficher.

accel0
accel1
accel2
accel3

Si aucun résultat ne s'affiche, cela signifie probablement que vous n'avez pas démarré le conteneur avec l'option --privileged.

Exécuter un modèle

Vous pouvez vérifier si l'appareil de VM TPU est disponible en exécutant la commande suivante:

(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