Exécuter le code PyTorch sur des tranches de pod TPU

PyTorch/XLA nécessite que toutes les VM TPU puissent accéder au code et aux données du modèle. Vous pouvez utiliser un script de démarrage pour télécharger le logiciel nécessaire à la distribution des données du modèle à toutes les VM TPU.

Si vous connectez vos VM TPU à un cloud privé virtuel (VPC), vous devez ajouter une règle de pare-feu dans votre projet pour autoriser l'entrée des ports 8470 à 8479. Pour en savoir plus sur l'ajout de règles de pare-feu, consultez Utiliser des règles de pare-feu

Configurer votre environnement

1. Dans Cloud Shell, exécutez la commande suivante pour vérifier que vous êtes exécutant la version actuelle de gcloud:

   $ gcloud components update

   Si vous devez installer gcloud, exécutez la commande suivante:

   $ sudo apt install -y google-cloud-sdk

2. Créez des variables d'environnement :

   $ export PROJECT_ID=project-id
   $ export TPU_NAME=tpu-name
   $ export ZONE=us-central2-b
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ export ACCELERATOR_TYPE=v4-32

Créez la VM TPU.

$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone=${ZONE} \
--project=${PROJECT_ID} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version ${RUNTIME_VERSION}

Configurer et exécuter le script d'entraînement

1. Ajoutez votre certificat SSH à votre projet :

   ssh-add ~/.ssh/google_compute_engine

2. Installer PyTorch/XLA sur tous les nœuds de calcul de VM TPU

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all --command="
     pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

3. Cloner XLA sur tous les nœuds de travail de VM TPU

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all --command="git clone -b r2.4 https://github.com/pytorch/xla.git"

4. Exécuter le script d'entraînement sur tous les nœuds de calcul

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all \
     --command="PJRT_DEVICE=TPU python3 ~/xla/test/test_train_mp_imagenet.py  \
     --fake_data \
     --model=resnet50  \
     --num_epochs=1 2>&1 | tee ~/logs.txt"
     

   L'entraînement prend environ cinq 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%
   

Effectuer un nettoyage

Lorsque vous avez fini d'utiliser votre VM TPU, procédez comme suit pour nettoyer vos ressources.

1. Déconnectez-vous de l'instance Compute Engine, si vous ne l'avez pas déjà fait :

   (vm)$ exit

   Votre invite devrait maintenant être username@projectname, indiquant que vous êtes dans Cloud Shell.

2. Supprimez vos ressources Cloud TPU et Compute Engine.

   $ gcloud compute tpus tpu-vm delete  \
     --zone=${ZONE}

3. Vérifiez que les ressources ont été supprimées en exécutant la commande gcloud compute tpus execution-groups list. La suppression peut prendre plusieurs minutes. Le résultat de la commande suivante ne doit inclure aucune des ressources créées dans ce tutoriel :

   $ gcloud compute tpus tpu-vm list --zone=${ZONE}