Exécuter 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 découvrir plus en détail comment entraîner un modèle sur un Cloud TPU, consultez l'un des Tutoriels PyTorch sur Cloud TPU

Avant de commencer

Avant de suivre ce guide de démarrage rapide, vous devez créer une instance Google Cloud Platform installez la Google Cloud CLI. et configurez la commande gcloud. Pour en savoir plus, consultez Configurez 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écuter:

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

Description des options de commande

zone
Zone où vous prévoyez de créer votre 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 Versions de TPU.
version
Logiciel Cloud TPU version.

Lors de la création de votre TPU, vous pouvez transmettre les --network supplémentaires et les options --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 le paramètre --network. L'option --subnetwork est facultative et peut servir à spécifiez un sous-réseau par défaut pour le réseau que vous utilisez (par défaut ou spécifié par l'utilisateur). Voir le gcloud page de la documentation de référence de l'API pour en savoir plus sur ces options.

Se connecter à la VM Cloud TPU

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

Installer PyTorch/XLA sur votre VM TPU

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

Définir la configuration de l'environnement d'exécution du 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-central2-b

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

Étape suivante

En savoir plus sur les VM Cloud TPU :