Cette page a été traduite par l'API Cloud Translation.

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 le guide de l'utilisateur de la VM Cloud TPU.

Consultez la page Versions compatibles de PyTorch pour obtenir la liste des versions du logiciel TPU disponibles pour Pytorch/XLA.

Configuration de base

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

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

   (vm)$ export TPU_NUM_DEVICES=4

Pour les modèles comportant 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 la VM TPU à utiliser la version standard de 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 version que vous souhaitez utiliser. Par exemple, si vous souhaitez utiliser 1.13:

(tpuvm):$ cd /usr/share/
(tpuvm):$ sudo git clone -b release/1.13 --recursive https://github.com/pytorch/pytorch
(tpuvm):$ cd pytorch/
(tpuvm):$ sudo git clone -b r1.13 --recursive https://github.com/pytorch/xla.git
(tpuvm):$ cd xla/
(tpuvm):$ yes | sudo pip3 uninstall torch_xla
(tpuvm):$ yes | sudo pip3 uninstall torch
(tpuvm):$ yes | sudo pip3 uninstall torch_vision
(tpuvm):$ sudo pip3 install torch==1.13.0
(tpuvm):$ sudo pip3 install torchvision==0.14.0
(tpuvm):$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.13-cp38-cp38-linux_x86_64.whl
(tpuvm):$ sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
(tpuvm):$ sudo pip3 install torch_xla[tpuvm]

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 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 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 dans 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 est conservé d'un cycle d'entraînement à un autre. 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 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.

Exemples de pods TPU PyTorch/XLA

Consultez la page Pod de VM TPU PyTorch pour obtenir des informations de configuration et des exemples d'exécution de PyTorch/XLA sur un pod de VM TPU.

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 fichier README de GitHub pour trouver toutes les images Docker de VM TPU disponibles.

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 la bibliothèque libtpu correspondant à l'aide de la commande suivante:

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

Vérifier `tcmalloc`

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

(pytorch) root:/# echo $LD_PRELOAD