Guide de démarrage rapide: exécuter un calcul sur une VM Cloud TPU à l'aide de PyTorch

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 la CLI Google Cloud et configurer la commande gcloud. Pour plus d'informations, consultez la page Configurer un compte et un projet Cloud TPU.

Créer un Cloud TPU avec gcloud

Lancez une VM Compute Engine et Cloud TPU à l'aide de la commande gcloud. La commande à utiliser varie selon que vous utilisez une VM TPU ou un nœud TPU. Pour en savoir plus sur les deux architectures de VM, consultez la page Architecture du système. Pour en savoir plus sur la commande gcloud, consultez la documentation de référence de gcloud.

VM TPU

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

Description des options de commande

zone
La zone dans laquelle vous prévoyez de créer votre Cloud TPU.
accelerator-type
Type du Cloud TPU à créer.
version
Version du logiciel Cloud TPU.

Nœud TPU

  1. Lors de la création d'un nœud TPU pour PyTorch, vous devez d'abord créer une instance de VM Compute Engine.

    gcloud compute instances create tpu-name \
    --zone=us-central1-b  \
    --machine-type=n1-standard-16  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=200GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform

    Description des options de commande

    project
    L'ID de votre projet GCP
    name
    Nom du Cloud TPU à créer.
    zone
    La zone dans laquelle vous prévoyez de créer votre Cloud TPU.
    disk-size
    Taille du disque dur en Go de la VM créée par la commande gcloud.
    machine-type
    Type de machine de la VM Compute Engine à créer.
    tf-version
    Version de TensorFlow gcloud installée sur la VM.
    accelerator-type
    Type du Cloud TPU à créer.
  2. Créez ensuite l'instance TPU.

    gcloud compute tpus create tpu-name \
    --zone=us-central1-b \
    --network=default \
    --version=pytorch-1.11 \
    --accelerator-type=v3-8

    Description des options de commande

    name
    Nom du Cloud TPU à créer.
    zone
    Zone dans laquelle vous prévoyez de créer votre Cloud TPU.
    network
    Réseau auquel le TPU est connecté.
    version
    Version de PyTorch gcloud installée sur la VM.
    accelerator-type
    Type du Cloud TPU à créer.

Se connecter à la VM Cloud TPU

VM TPU

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

Nœud TPU

gcloud compute ssh tpu-name --zone=us-central1-b

Définir la configuration de l'appareil TPU XRT

VM TPU

Configurez l'environnement Torch-XLA.

(vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

Nœud TPU

  1. Recherchez l'adresse IP du nœud TPU.

    (vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b
  2. Configurez l'environnement Torch-XLA. Veillez à remplacer your-tpu-ip-address par l'adresse IP de votre TPU.

    (vm)$ conda activate torch-xla-1.11
    (vm)$ export TPU_IP_ADDRESS=your-tpu-ip-address
    (vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

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')
    
    OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
    OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
    

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.

    VM TPU

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

    Nœud TPU

    $ gcloud compute tpus execution-groups 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 :