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
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.
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
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
Encontre o endereço IP do nó da TPU.
(vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b
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:
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)
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.
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.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: