Guia de início rápido do Cloud TPU 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 o SDK do Google Cloud Platform. e configure 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 uma VM ou um nó de TPU. Para ver 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 alpha compute tpus tpu-vm create tpu-name \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=v2-alpha

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 ambiente de execução do 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 do 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 gcloud do Tensorflow é 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.9 \
    --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 de PyTorch gcloud é instalada na VM.
    accelerator-type
    O tipo da Cloud TPU a ser criada.

Conecte-se à VM da Cloud TPU

VM de TPU

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

Nó da TPU

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

Definir 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ó de 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.9
    (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 da computação:

    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
    

Limpar

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

  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 alpha 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 da Cloud TPU: