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:
- Dernière image disponible: tpu-vm-pt-1.11
- tpu-vm-pt-2.10
- 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
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 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.
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=tpu-vm-pt-1.11
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.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écutantls -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 commandels
, 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 dexla_dist
. Les nouveaux paramètres du serveur XRT, tels que les variables d'environnement telles queLD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4
, ne prennent 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 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'
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
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 deusername@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.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%
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