Guide de l'utilisateur de Cloud TPU PyTorch/XLA

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.

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)
    

    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

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.

  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'
    

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

    (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é.

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

(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:/#

Valider libtpu

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