Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.
Guia de início rápido: executar um cálculo em uma VM da Cloud TPU usando o PyTorch

Executar um cálculo em uma VM do Cloud TPU usando o 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, crie uma conta do Google Cloud Platform, instale a Google Cloud CLI 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 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

Para criar uma VM de TPU no projeto de usuário padrão, execução de rede e computação/zona:

$ 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 do Cloud TPU a ser criado.
version
A versão do software da Cloud TPU.

Ao criar a TPU, é possível transmitir as sinalizações --network e --subnetwork adicionais caso queira especificar a rede e sub-rede padrão. Se você não quiser usar a rede padrão, será necessário transmitir a sinalização --network. A sinalização --subnetwork é opcional e pode ser usada para especificar uma sub-rede padrão para qualquer rede que você esteja usando (padrão ou especificada pelo usuário). Consulte a gcloud página de referência da API para saber detalhes sobre essas sinalizações.

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
    ID do projeto do Cloud
    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 do Cloud TPU a ser criado.

  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 do Cloud TPU a ser criado.

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 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 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: