Guide utilisateur de PyTorch/XLA Cloud TPU

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 projet ssh dans la procédure décrite dans la section Préparer un projet GCP. Les ressources nécessaires à l'exécution des commandes de ce guide sont 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 dotés d'allocations groupées de taille adaptée, tcmalloc améliore considérablement les performances par rapport à la valeur par défaut malloc. Ainsi, le malloc par défaut utilisé sur la VM TPU est tcmalloc. Toutefois, en fonction de votre charge de travail (par exemple, avec DLRM qui présente des allocations très élevées pour ses tables de représentations vectorielles continues), tcmalloc peut entraîner un ralentissement, auquel cas vous pouvez essayer de réinitialiser la variable suivante. en utilisant la valeur malloc par défaut à la place:

   (vm)$ unset LD_PRELOAD

Effectuez un calcul simple

Démarrez l'interpréteur Python 3:

(vm)$ python3

   Python 3.6.9 (default, Jan 26 2021, 15:33:00)
   [GCC 8.4.0] on linux
   Type "help", "copyright", "credits" or "license" for more information.

Depuis l'interpréteur Pytorch 3, importez les packages PyTorch suivants:

import torch
import torch_xla.core.xla_model as xm

Effectuez les calculs suivants:

dev = xm.xla_device()

Le résultat de cette commande ressemble à ce qui suit:

2021-04-01 23:20:23.268115: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
2021-04-01 23:20:23.269345: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey

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 tout code PyTorch / XLA si vous le souhaitez. Par exemple, vous pouvez exécuter une ResNet avec des données fictives à l'aide de la commande suivante:

   (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 s'entraîne pour une époque et dure environ 7 minutes. Un résultat semblable au suivant s'affiche:

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

Vérifiez que les ressources ont été supprimées en exécutant gcloud list. La suppression peut prendre plusieurs minutes. Si vous obtenez une réponse semblable à celle présentée ci-dessous, vos instances ont bien été supprimées.

   $ gcloud alpha compute tpus list --zone=zone

   Listed 0 items.

Configuration avancée

Dans les exemples ci-dessus (calcul simple et ResNet50), votre programme PyTorch/XLA démarrera le serveur XRT local dans le même processus. 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 dans l'entraînement. Dans ce cas, la journalisation côté serveur se trouve sous /tmp/xrt_server_log.

(vm)$ ls /tmp/xrt_server_log/

server_20210401-031010.log

Profilage des performances des VM TPU

Pour plus d'informations sur le profilage de vos modèles sur une VM TPU, consultez la page Profiler des performances PyTorch XLA.

Pods

PyTorch/XLA doit posséder toutes les VM TPU pour accéder au code et aux données du modèle. Pour ce faire, vous pouvez utiliser le script de démarrage suivant lors de la création du pod de VM TPU. Les données sont téléchargées sur toutes les VM TPU.

Dans Cloud Shell, exécutez la commande suivante pour vous assurer que vous utilisez la version la plus récente de gcloud:
```
$ sudo /opt/google-cloud-sdk/bin/gcloud components update
```

Exportez les variables d'environnement suivantes pour créer la VM TPU Pod:

$ 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 appliquez ces instructions, exécutez chaque commande commençant par (vm)$ dans la fenêtre de session de la VM.

ssh à n'importe quel nœud de calcul TPU, par exemple, nœud de calcul 0, vérifiez si le téléchargement des données/du modèle est terminé (cela prend une étape et démarrez l'entraînement après avoir généré les clés SSH vers ssh). entre les nœuds de calcul de la VM sur un pod.

Remarque :Pour les pods TPU, la valeur par défaut est ssh dans la première VM TPU. Pour ssh dans d'autres VM TPU associées au pod TPU, ajoutez --worker ${WORKER_NUMBER} dans la commande ssh, où l'index WORKER_NUMBER est 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 du projet ssh:
```
(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 pouvez trouver 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(au cas où le serveur était resté défaillant), 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 tels que LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 ne seront pas récupérés tant que vous n'aurez pas redémarré le 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 gcloud list. La suppression peut prendre plusieurs minutes. Si vous obtenez une réponse semblable à celle présentée ci-dessous, vos instances ont bien été supprimées.
```
$ gcloud alpha compute tpus list --zone=zone
```
```
Listed 0 items.
```

Pods à distance

Nous vous recommandons de configurer un coordinateur distant pour l'utiliser pour démarrer l'entraînement du pod. Le coordinateur distant est une VM Compute Engine standard, et non une VM TPU. Le coordinateur distant envoie des commandes au pod des VM TPU. L'avantage d'utiliser un coordinateur distant est qu'en cas d'événements de maintenance TPU, l'entraînement peut s'interrompre et vous risquez de perdre l'accès au pod TPU de la VM. Toutefois, le coordinateur distant peut autoriser la récupération automatique dans la VM TPU. s'affiche en haut de l'écran.

Exportez les variables d'environnement suivantes pour créer la VM TPU:

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

Créer une tranche de pod de VM TPU

Vous créez une tranche de pod TPU à l'aide de la commande gcloud. Par exemple, pour créer une tranche de pod v2-32, utilisez la commande suivante:

$ 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 pour créer la VM du coordinateur à distance:
```
$ export VM_NAME=vm-name
$ export ZONE=zone
```
Créez une VM de coordinateur à distance 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 et exécutez l'entraînement depuis cet emplacement.

(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 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 locale 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