Exécuter un calcul sur une VM Cloud TPU à l'aide de PyTorch

Ce document présente brièvement l'utilisation de PyTorch et de Cloud TPU.

Avant de commencer

Avant d'exécuter les commandes de ce document, vous devez créer un compte Google Cloud , installer Google Cloud CLI et configurer la commande gcloud. Pour en savoir plus, consultez la section Configurer l'environnement Cloud TPU.

Créer un Cloud TPU à l'aide de gcloud

  1. Définissez des variables d'environnement pour faciliter l'utilisation des commandes.

    export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name

    Descriptions des variables d'environnement

    PROJECT_ID
    L'ID de votre Google Cloud projet.
    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.
    ZONE
    Zone dans laquelle vous prévoyez de créer votre Cloud TPU.
    RUNTIME_VERSION
    Version du runtime Cloud TPU.
    TPU_NAME
    Nom attribué par l'utilisateur à votre Cloud TPU.

  2. Créez votre VM TPU en exécutant la commande suivante:

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

Se connecter à la VM Cloud TPU

Connectez-vous à votre VM TPU via SSH à l'aide de la commande suivante:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE

Installer PyTorch/XLA sur votre VM TPU

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html

Vérifier que PyTorch peut accéder aux TPU

Utilisez la commande suivante pour vérifier que PyTorch peut accéder à vos TPU.

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

Le résultat de la commande devrait ressembler à ceci:

['xla:0', 'xla:1', 'xla:2', 'xla:3']

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)$ PJRT_DEVICE=TPU 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 ne 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 \
  --project=$PROJECT_ID \
  --zone=$ZONE

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

Étape suivante

En savoir plus sur les VM Cloud TPU :