Criar uma fração de TPU com vários hosts

Criar uma fração de TPU com vários hosts

Saiba como criar uma fração de TPU de vários hosts usando um grupo gerenciado de instâncias (MIG), conectar-se à fração e executar um cálculo. Este guia de início rápido usa a opção de consumo sob demanda. Execute os comandos deste guia de início rápido no terminal local ou no Cloud Shell.

Antes de começar

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.

Instale a CLI do Google Cloud.

Observação:se você já instalou a CLI gcloud, execute o comando gcloud components update para verificar se tem a versão mais recente.

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

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

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.

Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

• 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 .

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.

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

Ative a API Compute Engine:

Funções necessárias para ativar APIs

Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão com o papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão com o papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

gcloud services enable compute.googleapis.com

Instale a CLI do Google Cloud.

Observação:se você já instalou a CLI gcloud, execute o comando gcloud components update para verificar se tem a versão mais recente.

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

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

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.

Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

• 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 .

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.

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

Ative a API Compute Engine:

Funções necessárias para ativar APIs

Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão com o papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão com o papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

gcloud services enable compute.googleapis.com

Funções exigidas

Para ter as permissões necessárias para criar um MIG que forma uma fração de TPU de vários hosts, conecte-se a cada VM no MIG usando SSH e execute comandos, peça ao administrador para conceder a você os seguintes papéis do IAM no projeto:

• Administrador da instância do Compute (v1) (roles/compute.instanceAdmin.v1)
• Usuário da conta de serviço (roles/iam.serviceAccountUser)
• Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin)

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 um modelo de instância

Para criar um modelo de instância para VMs de TPU v6e, use o comando gcloud compute instance-templates create:

Observação: se você usar um tipo de máquina diferente, verifique se está usando uma imagem de SO compatível. Para conferir uma lista de imagens de SO gerenciadas pelo Google para cada versão de TPU compatível, consulte Imagens de SO da TPU.

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

Criar uma política de carga de trabalho

Uma política de carga de trabalho define propriedades físicas das suas instâncias de computação. Em frações de TPU, a topologia do acelerador define a disposição física dos chips de TPU. É necessário especificar uma topologia de acelerador para frações de TPU interconectadas de vários hosts.

Para criar uma política de carga de trabalho para uma fração de TPU de vários hosts, use o comando gcloud compute resource-policies create workload-policy com a flag --accelerator-topology. O comando a seguir cria uma política de carga de trabalho com uma topologia 2x4:

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

Criar um MIG

Execute os comandos a seguir para criar uma MIG que forma uma fração de TPU de vários hosts.

Observação: talvez você encontre um erro temporário de disponibilidade de recursos enquanto o MIG aguarda na fila por capacidade. Se você tiver longos tempos de espera, tente outra região. Para conferir uma lista de locais disponíveis, consulte Regiões e zonas de TPU. Se você escolher uma região diferente, primeiro crie uma política de carga de trabalho nessa região repetindo a etapa Criar uma política de carga de trabalho. Atualize a região em todos os comandos subsequentes.

1. Para criar um MIG que forma uma fração de TPU de vários hosts, use o comando gcloud compute instance-groups managed create:

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   Substitua PROJECT_ID pelo ID do projeto Google Cloud .

2. Se quiser, verifique se as instâncias gerenciadas estão em execução usando os seguintes comandos:

   • Para conferir o status geral do MIG, use o comando gcloud compute instance-groups managed describe:

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • Para conferir o status de cada instância, use o comando gcloud compute instance-groups managed list-instances:

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

Instalar o JAX

Instale as dependências e o framework JAX em um ambiente virtual em todas as instâncias de VM de TPU no MIG. Se as VMs da TPU tiverem uma versão do Python anterior à 3.11 instalada, instale o Python 3.11 para executar a versão mais recente do JAX.

1. Verifique qual versão do Python está sendo executada nas suas VMs da TPU:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   Se a versão for anterior ao Python 3.11, instale o Python 3.11:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. Crie um ambiente virtual:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. Instale o JAX no ambiente virtual:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

Executar o código JAX na fração

Para executar o código JAX em uma fração de TPU, é preciso executá-lo em cada host dessa fração. A chamada de função jax.device_count() para de responder até ser chamada em cada host na fração. O exemplo a seguir mostra como executar um cálculo JAX em uma fração de TPU.

Preparar o código

Crie um arquivo chamado example.py em cada instância:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

Executar o código na fração

Execute o programa example.py em cada VM de TPU na fração:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

A saída será semelhante a esta:

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]

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.

Se quiser manter o projeto, exclua apenas o MIG e todas as VMs do grupo usando o comando gcloud compute instance-groups managed delete:

gcloud compute instance-groups managed delete quickstart-tpu-mig --region=us-east5

A seguir

• Saiba mais sobre os recursos de TPU no Compute Engine
• Criar uma fração de TPU com vários hosts usando outras opções de consumo
• Saiba mais sobre instâncias de TPU em MIGs.