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 accéder à un tutoriel plus détaillé illustrant l'entraînement d'un modèle sur un Cloud TPU, consultez l'un des tutoriels Cloud TPU PyTorch.

Avant de commencer

Avant de suivre ce guide de démarrage rapide, vous devez créer un compte Google Cloud Platform, installer 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, le réseau et la zone utilisateur par défaut, exécutez les commandes suivantes:

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

Description des options de commande

zone
Zone dans laquelle vous prévoyez de créer votre Cloud TPU.
accelerator-type
Le type d'accélérateur spécifie la version et la taille du 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.

Lorsque vous créez votre TPU, vous pouvez transmettre les indicateurs --network et --subnetwork supplémentaires 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'indicateur --network. L'indicateur --subnetwork est facultatif et peut être utilisé 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 indicateurs, 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-a

Installer PyTorch/XLA sur votre VM TPU

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

Définir la configuration d'exécution du TPU

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

   (vm) $ export PJRT_DEVICE=TPU

Vérifier que PyTorch peut accéder aux TPU

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

    import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')

  2. Exécutez le script :

    (vm)$ python3 tpu-count.py

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

    PyTorch can access 8 TPU cores

Effectuer un calcul de base

  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-a

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

Étape suivante

En savoir plus sur les VM Cloud TPU :