Guía de inicio rápido de Cloud TPU PyTorch

En esta guía de inicio rápido, se muestra cómo crear una Cloud TPU, instalar PyTorch y ejecutar un cálculo simple en una Cloud TPU. Para obtener un instructivo más detallado en el que se muestra cómo entrenar un modelo en Cloud TPU, consulta uno de los instructivos de Cloud TPU PyTorch.

Antes de comenzar

Antes de seguir esta guía de inicio rápido, debes crear una cuenta de Google Cloud Platform y, luego, instalar el SDK de Google Cloud Platform y configurar el comando de gcloud. Para obtener más información, consulta Configura una cuenta y un proyecto de Cloud TPU.

Crea una Cloud TPU con gcloud

Inicia una VM de Compute Engine y Cloud TPU con el comando gcloud. El comando que uses depende de si usas una VM de TPU o un nodo TPU. Para obtener más información sobre la arquitectura de las dos VM, consulta Arquitectura del sistema. Para obtener más información sobre el comando de gcloud, consulta la Referencia de gcloud.

VM de TPU

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

Descripciones de las marcas de comandos

zone
Es la zona en la que deseas crear la Cloud TPU.
accelerator-type
El tipo de Cloud TPU que se creará.
version
La versión del entorno de ejecución de Cloud TPU.

Nodo TPU

  1. Cuando creas un nodo TPU para PyTorch, primero debes crear una instancia de VM de 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

    Descripciones de las marcas de comandos

    project
    : El ID de tu proyecto de GCP
    name
    El nombre de la Cloud TPU para crear.
    zone
    Es la zona en la que deseas crear la Cloud TPU.
    disk-size
    El tamaño del disco duro en GB de la VM creada con el comando gcloud.
    machine-type
    El tipo de máquina de la VM de Compute Engine que se creará.
    tf-version
    La versión de Tensorflow gcloud se instala en la VM.
    accelerator-type
    El tipo de Cloud TPU que se creará.
  2. A continuación, crea la instancia de TPU.

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

    Descripciones de las marcas de comandos

    name
    El nombre de la Cloud TPU para crear.
    zone
    Es la zona en la que deseas crear la Cloud TPU.
    network
    La red a la que está conectada la TPU.
    version
    La versión de Pytorch gcloud se instala en la VM.
    accelerator-type
    El tipo de Cloud TPU que se creará.

Conéctate a tu VM de Cloud TPU

VM de TPU

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

Nodo TPU

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

Establecer la configuración del dispositivo TPU XRT

VM de TPU

Configura el entorno Torch-XLA.

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

Nodo TPU

  1. Busca la dirección IP del nodo TPU.

    (vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b
  2. Configura el entorno Torch-XLA. Asegúrate de reemplazar your-tpu-ip-address por la dirección IP de tu TPU.

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

Realiza un cálculo simple:

  1. Crea un archivo llamado tpu-test.py en el directorio actual y copia y pega la siguiente secuencia de comandos en él.

    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. Ejecuta la secuencia de comandos:

      (vm)$ python3 tpu-test.py

    En el resultado de la secuencia de comandos, se muestra el resultado del cálculo:

    tensor([[-0.2121,  1.5589, -0.6951],
            [-0.7886, -0.2022,  0.9242],
            [ 0.8555, -1.8698,  1.4333]], device='xla:1')</span></pre>
    
    OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
    OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
    

Limpia

Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta página.

  1. Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

    (vm)$ exit
    

    El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

  2. Borra tu Cloud TPU.

    VM de TPU

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

    Nodo TPU

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

    El resultado de este comando debe confirmar que la TPU se haya borrado.

¿Qué sigue?

Obtén más información sobre las VM de Cloud TPU: