Criar uma instância de TPU

Neste documento, mostramos como criar uma instância de máquina virtual (VM) da TPU, conectar-se a ela e executar um cálculo nela. A instância de TPU criada neste guia de início rápido usa a opção on demand, que é a opção de consumo padrão.

Antes de começar

  1. Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

  2. Instale a CLI do Google Cloud.

  3. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

  4. Para inicializar a CLI gcloud, execute o seguinte comando: 

    gcloud init

  5. Crie ou selecione um Google Cloud projeto.

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  6. Se este guia estiver usando um projeto atual, verifique se você tem as permissões necessárias para concluir o guia. Se você criou um projeto, já tem as permissões necessárias.

  7. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  8. Ative a API Compute Engine:

    Funções necessárias para ativar APIs

    Para ativar as APIs, é necessário ter o papel do IAM de administrador de uso do serviço (roles/serviceusage.serviceUsageAdmin), que contém a permissão serviceusage.services.enable. Saiba como conceder papéis. 

    gcloud services enable compute.googleapis.com

  9. Instale a CLI do Google Cloud.

  10. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

  11. Para inicializar a CLI gcloud, execute o seguinte comando: 

    gcloud init

  12. Crie ou selecione um Google Cloud projeto.

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  13. Se este guia estiver usando um projeto atual, verifique se você tem as permissões necessárias para concluir o guia. Se você criou um projeto, já tem as permissões necessárias.

  14. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  15. Ative a API Compute Engine:

    Funções necessárias para ativar APIs

    Para ativar as APIs, é necessário ter o papel do IAM de administrador de uso do serviço (roles/serviceusage.serviceUsageAdmin), que contém a permissão serviceusage.services.enable. Saiba como conceder papéis. 

    gcloud services enable compute.googleapis.com
    16.

Funções exigidas

Para receber as permissões necessárias para criar uma instância de TPU, se conectar a ela usando SSH e executar comandos, peça ao administrador para conceder a você os seguintes papéis do IAM no projeto:

Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

Também é possível conseguir as permissões necessárias usando papéis personalizados ou outros papéis predefinidos.

Criar uma instância de TPU

Execute os comandos a seguir no terminal local ou usando o Cloud Shell.

  1. Para criar uma instância de TPU, execute o seguinte comando gcloud compute instances create:

    gcloud compute instances create quickstart-tpu \
    --machine-type=ct6e-standard-4t \
    --zone=us-east5-b \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accelerator-2204-amd64-with-tpu-v6e

  2. Se quiser, verifique se a instância está em execução usando o comando gcloud compute instances describe:

    gcloud compute instances describe quickstart-tpu --zone=us-east5-b

    Na saída, verifique o campo status para verificar se a instância está em execução ou ainda pendente.

Conectar-se à instância de TPU

É possível se conectar a uma instância de TPU usando o console do Google Cloud ou a CLI gcloud.

Para se conectar à VM usando SSH, execute o seguinte comando gcloud compute ssh:

gcloud compute ssh quickstart-tpu --project=PROJECT_ID --zone=us-east5-b

Instalar o JAX

Instale o framework JAX executando o seguinte comando na instância de TPU:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Executar uma verificação do sistema

Verifique se o JAX pode acessar a TPU e executar operações básicas com os seguintes comandos na instância de TPU:

  1. Inicie o interpretador do Python 3:

    python3

  2. Importe o módulo JAX:

    >>> import jax

  3. Confira o número de núcleos de TPU disponíveis:

    >>> jax.device_count()

    A saída mostra o número de núcleos de TPU. O número de núcleos depende da versão da TPU que você está usando. Para mais informações, consulte Versões de TPU.

Fazer um cálculo

  1. Execute o comando a seguir para somar dois números:

    >>> jax.numpy.add(1, 1)

    A saída será semelhante a esta:

    >>> Array(2, dtype=int32, weak_type=True)

  2. Saia do interpretador Python:

    >>> exit()

Limpar

Para evitar cobranças na conta do Google Cloud pelos recursos usados nesta página, exclua o projeto do Google Cloud e os recursos.

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados nesta página, exclua a instância de TPU usando o comando gcloud compute instances delete:

gcloud compute instances delete quickstart-tpu --zone=us-east5-b

A seguir