Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.
Guide de démarrage rapide: effectuer 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 Google Cloud CLI 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

Pour créer une VM TPU dans le projet utilisateur par défaut, le réseau et l'exécution de calcul/zone:

$ 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
Zone dans laquelle vous prévoyez de créer votre Cloud TPU.
accelerator-type
Le type de Cloud TPU à créer.
version
Version du logiciel Cloud TPU.

Lors de la création de votre TPU, vous pouvez transmettre les options --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'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 indicateurs, consultez la page de référence de l'API gcloud.

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
    Votre ID de projet Cloud
    name
    Nom du Cloud TPU à créer.
    zone
    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
    La version de TensorFlow gcloud est installée sur la VM.
    accelerator-type
    Le type de 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
    La version de PyTorch gcloud est installée sur la VM.
    accelerator-type
    Le type de 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

Afin d'éviter que des frais ne soient facturés sur votre compte Google Cloud pour les ressources utilisées sur cette page, 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 :