Effectuer un calcul sur une VM Cloud TPU à l'aide de PyTorch

Ce guide de démarrage rapide explique comment créer un Cloud TPU, installer PyTorch et exécuter un calcul simple sur Cloud TPU. Pour un tutoriel plus détaillé montrant comment entraîner un modèle sur Cloud TPU, consultez l'un des tutoriels Cloud TPU sur PyTorch.

Avant de commencer

Avant de suivre ce guide de démarrage rapide, vous devez créer un compte Google Cloud Platform, installer la Google Cloud CLI et configurer la commande gcloud. Pour en savoir plus, consultez la page Configurer un compte et un projet Cloud TPU.

Créer un Cloud TPU avec gcloud

Pour créer une VM TPU dans le projet utilisateur, le réseau et la zone par défaut, exécutez la commande suivante:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-b \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Description des options de commande

zone
Zone dans laquelle vous prévoyez de créer la ressource Cloud TPU.
accelerator-type
Le type d'accélérateur spécifie la version et la taille de la ressource Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
version
Version du logiciel Cloud TPU.

Lors de la création de votre TPU, vous pouvez transmettre les options supplémentaires --network et --subnetwork si vous souhaitez spécifier le réseau et le sous-réseau par défaut. Si vous ne souhaitez pas utiliser le réseau par défaut, vous devez transmettre l'option --network. L'option --subnetwork est facultative et peut être utilisée pour spécifier un sous-réseau par défaut pour le réseau que vous utilisez (par défaut ou spécifié par l'utilisateur). Pour en savoir plus sur ces options, consultez la page de référence de l'API gcloud.

Se connecter à la VM Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-b

Installer PyTorch/XLA sur votre VM TPU

   (vm)$ pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

Définir la configuration de l'environnement d'exécution TPU

Assurez-vous que l'environnement d'exécution PyTorch/XLA utilise le TPU.

   (vm) $ export PJRT_DEVICE=TPU

Effectuez un calcul simple :

  1. Créez un fichier nommé tpu-test.py dans le répertoire actuel et collez-y le script suivant.

    import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

  2. Exécutez le script :

      (vm)$ python3 tpu-test.py

    Le résultat du script montre le résultat du calcul :

    tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Effectuer un nettoyage

Pour éviter que les ressources utilisées sur cette page soient facturées sur votre compte Google Cloud, procédez comme suit :

  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 votre Cloud TPU.

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-b

Le résultat de cette commande doit confirmer que votre TPU a été supprimé.

Étapes suivantes

En savoir plus sur les VM Cloud TPU :