Guia de início rápido: como executar um cálculo em uma VM da Cloud TPU usando o PyTorch

Executar um cálculo em uma VM da Cloud TPU usando PyTorch

Neste guia de início rápido, mostramos como criar um Cloud TPU, instalar o PyTorch e executar um cálculo simples em um Cloud TPU. Para ver um tutorial mais detalhado mostrando como treinar um modelo em um Cloud TPU, consulte um dos Tutoriais do Cloud TPU PyTorch.

Antes de começar

Antes de seguir este guia de início rápido, é preciso criar uma conta do Google Cloud Platform, instalar a CLI do Google Cloud e configurar o comando gcloud. Para mais informações, consulte Configurar uma conta e um projeto do Cloud TPU.

Criar um Cloud TPU com gcloud

Use o comando gcloud para iniciar uma VM do Compute Engine e a Cloud TPU. O comando usado depende de você estar usando uma VM ou um nó de TPU. Para mais informações sobre a arquitetura de duas VMs, consulte Arquitetura do sistema. Para mais informações sobre o comando gcloud, consulte a Referência do gcloud.

VM de TPU

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

Descrições de sinalizações de comando

zone
A zona em que você planeja criar a Cloud TPU.
accelerator-type
O tipo da Cloud TPU a ser criada.
version
A versão do software da Cloud TPU.

Nó da TPU

  1. Ao criar um nó de TPU para PyTorch, primeiro crie uma instância de VM do 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

    Descrições de sinalizações de comando

    project
    O ID do projeto no GCP
    name
    O nome do Cloud TPU a ser criado.
    zone
    A zona em que você planeja criar a Cloud TPU.
    disk-size
    O tamanho do disco rígido em GB da VM criada pelo comando gcloud.
    machine-type
    O tipo de máquina da VM do Compute Engine a ser criada.
    tf-version
    A versão do TensorFlow gcloud é instalada na VM.
    accelerator-type
    O tipo da Cloud TPU a ser criada.
  2. Em seguida, crie a instância da TPU.

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

    Descrições de sinalizações de comando

    name
    O nome do Cloud TPU a ser criado.
    zone
    A zona em que você planeja criar a Cloud TPU.
    network
    A rede à qual a TPU está conectada.
    version
    A versão do PyTorch gcloud é instalada na VM.
    accelerator-type
    O tipo da Cloud TPU a ser criada.

Conecte-se à VM do Cloud TPU

VM de TPU

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

Nó da TPU

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

Defina a configuração do dispositivo XRT TPU

VM de TPU

Configure o ambiente Torch-XLA.

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

Nó da TPU

  1. Encontre o endereço IP do nó da TPU.

    (vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b
  2. Configure o ambiente Torch-XLA. Certifique-se de substituir your-tpu-ip-address pelo endereço IP da sua 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"

Faça um cálculo simples:

  1. Crie um arquivo chamado tpu-test.py no diretório atual e copie e cole o script a seguir nele.

    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. Execute o script:

      (vm)$ python3 tpu-test.py

    A saída do script mostra o resultado do cálculo:

    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
    

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados nesta página, siga as etapas abaixo.

  1. Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:

    (vm)$ exit
    

    Agora, o prompt precisa ser username@projectname, mostrando que você está no Cloud Shell.

  2. Exclua o Cloud TPU.

    VM de TPU

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

    Nó da TPU

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

    A saída desse comando confirmará se a TPU foi excluída.

A seguir

Leia mais sobre VMs do Cloud TPU: