Implantar cargas de trabalho de TPU no GKE Standard

Padrão

---

Nesta página, você vai aprender a solicitar e implantar cargas de trabalho que usam aceleradores de Cloud TPU (TPUs) no Google Kubernetes Engine (GKE).

Antes de configurar e implantar cargas de trabalho da TPU no GKE, é necessário estar familiarizado com os seguintes conceitos:

1. Introdução ao Cloud TPU.
2. Arquitetura do sistema do Cloud TPU.
3. Sobre TPUs no GKE.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

• Ativar a API Google Kubernetes Engine.
Ativar a API Google Kubernetes Engine
• Se você quiser usar a Google Cloud CLI para essa tarefa, instale e, em seguida, inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando gcloud components update.
  Observação: para instalações atuais da CLI gcloud, defina as propriedades compute/region e compute/zone. Quando você define os locais padrão, você evita erros como este na gcloud CLI: One of [--zone, --region] must be supplied: Please specify location.

Disponibilidade da TPU no GKE

Use o GKE para criar e gerenciar pools de nós com TPUs. É possível usar esses aceleradores criados intencionalmente para realizar treinamento, ajuste e inferência de modelos de IA em grande escala.

Confira uma lista de versões de TPU com suporte no GKE.

Planejar a configuração da TPU

Planeje a configuração de TPU com base no modelo de machine learning e na quantidade de memória necessária. Veja a seguir as etapas relevantes ao planejar a configuração da TPU:

1. Selecione uma versão e topologia de TPU.
2. Selecione o tipo de pool de nós a ser usado.

Garanta cota suficiente para VMs spot ou sob demanda

Se você estiver criando um pool de nós de TPU com VMs spot ou sob demanda, é preciso ter cota de TPU suficiente disponível na região que você quer usar.

A criação de um pool de nós de TPU que consome uma reserva de TPU não exige nenhuma cota de TPU.¹ É possível pular essa etapa com segurança para TPUs reservadas.

Criar um pool de nós de TPU sob demanda ou Spot no GKE requer cota da API Compute Engine. A cota da API Compute Engine (compute.googleapis.com) não é igual à cota da API Cloud TPU (tpu.googleapis.com), que é necessária para criar TPUs com a API Cloud TPU.

Para verificar o limite e o uso atual da sua cota da API Compute Engine para TPUs, siga estas etapas:

1. Acesse a página Cotas no console do Google Cloud.

   Acessar "Cotas"

2. Na caixa Filtro , faça o seguinte:

   1. Selecione a propriedade Serviço, insira API Compute Engine e pressione Enter.

   2. Selecione a propriedade Tipo e escolha Cota.

   3. Selecione a propriedade Nome e digite o nome da cota com base na versão da TPU e no tipo de máquina. Por exemplo, se você planeja criar nós da TPU v5e sob demanda com o tipo de máquina que começa com ct5lp-, digite TPU v5 Lite PodSlice chips.

      Versão da TPU	O tipo de máquina começa com	Nome da cota para instâncias sob demanda	Nome da cota para instâncias para Spot^2
      TPU v4	ct4p-	TPU v4 PodSlice chips	Preemptible TPU v4 PodSlice chips
      TPU v5e	ct5l-	TPU v5 Lite Device chips	Preemptible TPU v5 Lite Device chips
      TPU v5e	ct5lp-	TPU v5 Lite PodSlice chips	Preemptible TPU v5 Lite PodSlice chips
      TPU v5p	ct5p-	TPU v5p chips	Preemptible TPU v5p chips

   4. Selecione a propriedade Dimensões (por exemplo, locais) e insira region: seguido do nome da região em que você planeja criar TPUs no GKE. Por exemplo, insira region:us-west4 se planeja criar nós da TPU na zona us-west4-a. A cota de TPU é regional, portanto, todas as zonas na mesma região consomem a mesma cota de TPU.

Se nenhuma cota corresponder ao filtro inserido, isso significa que o projeto não recebeu nenhuma das cotas especificadas para a região desejada. Você precisará solicitar um aumento de cota de TPU.

Observação: quando uma reserva de TPU é criada, os valores de limite e de uso atual da cota correspondente aumentam pelo número de ícones na reserva da TPU. Por exemplo, quando uma reserva é criada para 16 chips de TPU v5e com tipo de máquina que começa com ct5lp-, o Limite e o Uso atual para a cota de TPU v5 Lite PodSlice chips na região relevante aumenta em 16.

1. Ao criar um pool de nós da TPU, use as flags --reservation e --reservation-affinity=specific para criar uma instância reservada. As reservas de TPU estão disponíveis ao comprar um compromisso. ↩

2. Ao criar um pool de nós de TPU, use a flag --spot para criar uma instância de Spot. ↩

Garantir a disponibilidade de reserva

A criação de um pool de nós de TPU reservado, ou seja, um pool de nós de TPU que consome uma reserva, não requer nenhuma cota de TPU. No entanto, a reserva precisa ter ícones suficientes disponíveis ou não utilizados no momento em que o pool de nós é criado.

Para ver quais reservas existem em um projeto, confira uma lista de reservas.

Para ver quantos ícones estão disponíveis em uma reserva de TPU, veja os detalhes de uma reserva.

criar um cluster

Crie um cluster do GKE no modo padrão em uma região com TPUs disponíveis. Recomendamos que você use clusters regionais, que fornecem alta disponibilidade do plano de controle do Kubernetes. Use a Google Cloud CLI ou o console do Google Cloud.

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

Substitua:

• CLUSTER_NAME: o nome do novo cluster;
• LOCATION: a região com sua capacidade de TPU disponível.
• VERSION: a versão do GKE, que precisa ser compatível com o tipo de máquina que você quer usar. A versão padrão do GKE pode não ter disponibilidade para sua TPU de destino. Para saber quais são as versões mínimas do GKE disponíveis por tipo de máquina de TPU, consulte Disponibilidade de TPU no GKE

Criar um pool de nós

Fatia de TPU de host único

É possível criar um pool de nós de fatia da TPU de host único usando a CLI do Google Cloud, o Terraform ou o console do Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    [--num-nodes=NUM_NODES \]
    [--spot \]
    [--enable-autoscaling \]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME \]
    [--total-min-nodes TOTAL_MIN_NODES \]
    [--total-max-nodes TOTAL_MAX_NODES \]
    [--location-policy=ANY]

Substitua:

• POOL_NAME: o nome do novo pool de nós.

• LOCATION: o nome da zona com base na versão da TPU que você quer usar:

  • Para a TPU v4, use us-central2-b.
  • Para os tipos de máquina TPU v5e que começam com ct5l-, use us-central1-a ou europe-west4-b.
  • Para os tipos de máquina TPU v5e que começam com ct5lp-, use us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, ou europe-west4-a.
  • Para a TPU v5p, use us-east1-d, us-east5-a ou us-east5-c.

  Para saber mais, consulte a disponibilidade da TPU no GKE.

• CLUSTER_NAME: o nome do cluster.

• NODE_ZONE: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.

• MACHINE_TYPE: o tipo de máquina a ser usado para nós. Para mais informações sobre os tipos de máquinas compatíveis com TPU, use a tabela em Mapeamento da configuração de TPU.

Também é possível usar as seguintes sinalizações:

• NUM_NODES: o número inicial de nós no pool em cada zona. Se você omitir essa sinalização, o padrão será 3. Se o escalonamento automático estiver ativado para o pool de nós usando a sinalização --enable-autoscaling, recomendamos que você defina NUM_NODES como 0, já que o escalonador automático provisiona nós adicionais assim que sua cargas de trabalho exige.
• RESERVATION_NAME: o nome da reserva que o GKE usa ao criar o pool de nós. Se você omitir essa sinalização, o GKE usará as TPUs disponíveis. Para saber mais sobre reservas de TPU, consulte Reserva de TPU.
• --enable-autoscaling: Adicionar um pool de nós com escalonamento automático ativado.
  • TOTAL_MIN_NODES: O número mínimo de nós no pool. Omita esse campo, a menos que o escalonamento automático também seja especificado.
  • TOTAL_MAX_NODES: O número máximo de nós no pool. Omita esse campo, a menos que o escalonamento automático também seja especificado.
• --spot: define o pool de nós para usar VMs spot para os nós no pool de nós. Isso não pode ser alterado após a criação do pool de nós.

Terraform

1. Use a versão 4.84.0 ou mais recente do provedor google.
2. Adicione o seguinte bloco à configuração do Terraform:

resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES
  autoscaling {
    total_min_node_count = TOTAL_MIN_NODES
    total_max_node_count = TOTAL_MAX_NODES
    location_policy      = "ANY"
  }

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }
}

Substitua:

• NODE_POOL_RESOURCE_NAME: o nome do recurso do pool de nós no modelo do Terraform.
• PROJECT_ID: o ID do projeto.
• CLUSTER_NAME: o nome do cluster existente.
• POOL_NAME: o nome do pool de nós a ser criado.
• CLUSTER_LOCATION: as zonas de computação do cluster. Especifique a região em que a versão da TPU está disponível. Para saber mais, consulte Selecionar uma versão de TPU e topologia.
• NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.
• NUM_NODES: o número inicial de nós no pool em cada uma das zonas do pool de nós. Se omitido, o padrão é 3. Se o escalonamento automático estiver ativado para o pool de nós usando o modelo de austoscaling, recomendamos que você defina NUM_NODES para 0, já que o GKE provisiona nós de TPU adicionais assim que sua carga de trabalho exige.
• MACHINE_TYPE: o tipo de máquina da TPU a ser usada. Para ver os tipos de máquina compatíveis com a TPU, use a tabela em Mapeamento da configuração da TPU.

Também é possível usar as seguintes variáveis:

• autoscaling: Adicionar um pool de nós com escalonamento automático ativado. Para uma fração de TPU de host único, o GKE faz o escalonamento entre os valores TOTAL_MIN_NODES e TOTAL_MAX_NODES.
  • TOTAL_MIN_NODES: O número mínimo de nós no pool. Esse campo é opcional, a menos que o escalonamento automático também seja especificado.
  • TOTAL_MAX_NODES: o número máximo de todos os nós no pool. Esse campo é opcional, a menos que o escalonamento automático também seja especificado.
• RESERVATION_NAME: se você usar a reserva de TPU, esta será a lista de rótulos dos recursos de reserva a serem usados ao criar o pool de nós. Para saber como preencher o RESERVATION_LABEL_VALUES no campo reservation_affinity, consulte Provedor do Terraform.
• spot: define o pool de nós para usar VMs do Spot nos nós da TPU. Isso não pode ser alterado após a criação do pool de nós. Para mais informações, consulte VMs spot.

Console

Para criar um pool de nós com TPUs:

1. Acesse a página Google Kubernetes Engine no console do Google Cloud.

   Acessar o Google Kubernetes Engine

2. Na lista de clusters, clique no nome do cluster que você quer modificar.

3. Clique em Adicionar pool de nós.

4. Na seção Detalhes do pool de nós, marque a caixa Especificar locais do nó.

5. Selecione a zona com base na versão de TPU que você quer usar:

   • Para a TPU v4, use us-central2-b.
   • Para os tipos de máquina TPU v5e que começam com ct5l-, use us-central1-a ou europe-west4-b.
   • Para os tipos de máquina TPU v5e que começam com ct5lp-, use us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, ou europe-west4-a.
   • Para a TPU v5p, use us-east1-d, us-east5-a ou us-east5-c.

6. No painel de navegação, clique em Nós.

7. Na seção Configuração da máquina, selecione TPUs.

8. No menu suspenso Série, selecione uma das seguintes opções:

   • CT4P: TPU v4.
   • CT5LP: TPU v5e
   • CT5P: TPU v5p

9. No menu suspenso Tipo de máquina, selecione o nome da máquina que será usada para os nós. Use a tabela Mapeamento da configuração de TPU para saber como definir o tipo de máquina e a topologia de TPU que criam um pool de nós de TPU de host único.

10. No menu suspenso Topologia da TPU, selecione a topologia física para a fatia da TPU.

11. Na caixa de diálogo Alterações necessárias, clique em Fazer alterações.

12. Certifique-se de que o Tipo de disco de inicialização seja Disco permanente padrão ou Disco permanente SSD.

13. Se quiser, marque a caixa de seleção Ativar nós em VMs do Spot para usar as VMs do Spot nos nós do pool.

14. Clique em Criar.

Fatia de TPU de vários hosts

É possível criar um pool de nós em uma fração da TPU de vários hosts usando a CLI do Google Cloud, o Terraform ou o console do Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

Substitua:

• POOL_NAME: o nome do novo pool de nós.

• LOCATION: o nome da zona com base na versão da TPU que você quer usar:

  • Para a TPU v4, use us-central2-b.
  • Os tipos de máquina TPU v5e que começam com ct5l- nunca são de vários hosts.
  • Para os tipos de máquina TPU v5e que começam com ct5lp-, use us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, ou europe-west4-a.
  • Para os tipos de máquina TPU v5p que começam com ct5p-, use us-east1-d, us-east5-a ou us-east5-c.

  Para saber mais, consulte a disponibilidade da TPU no GKE.

• CLUSTER_NAME: o nome do cluster.

• NODE_ZONE: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.

• MACHINE_TYPE: o tipo de máquina a ser usado para nós. Para saber mais sobre os tipos de máquina disponíveis, consulte Mapeamento da configuração de TPU.

• TPU_TOPOLOGY: a topologia física da fatia de TPU. O formato da topologia depende da versão da TPU, da seguinte maneira:

  • TPU v4 ou v5p: defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4.
  • TPU v5e: defina a topologia em duas tuplas ({A}x{B}), por exemplo, 2x2.

  Para saber mais, consulte Topologia.

• NUM_NODES: o número de nós no pool de nós. Precisa ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada VM. Para a TPU v4e de vários hosts e a TPU v5e, o número de chips em cada VM é quatro. Portanto, se a TPU_TOPOLOGY for 2x4x4 (TPU v4 com quatro chips em cada VM), a NUM_NODES será 32/4, que é igual a 8.

Opcionalmente, é possível usar as seguintes sinalizações:

• RESERVATION_NAME: o nome da reserva que o GKE usa ao criar o pool de nós. Se você omitir essa sinalização, o GKE usará pools de nós da TPU disponíveis. Para saber mais sobre reservas de TPU, consulte Reserva de TPU.
• --spot: define o pool de nós para usar VMs do Spot nos nós da TPU. Isso não pode ser alterado após a criação do pool de nós. Para mais informações, consulte VMs spot.
• --enable-autoscaling: adicionar um pool de nós com escalonamento automático ativado. Quando o GKE escalona um pool de nós de uma fração da TPU de vários hosts, ele escalona horizontalmente e de maneira atomizada esse pool de nós de zero até o tamanho máximo.
  • MAX_NODES: o tamanho máximo do pool de nós. A flag --max-nodes será obrigatória se --enable-autoscaling for fornecido e precisar ser igual ao produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de ícones para cada VM.

Terraform

1. Use a versão 4.84.0 ou mais recente do provedor google.

2. Adicione o seguinte bloco à configuração do Terraform:

   resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
     provider           = google
     project            = PROJECT_ID
     cluster            = CLUSTER_NAME
     name               = POOL_NAME
     location           = CLUSTER_LOCATION
     node_locations     = [NODE_ZONES]
     initial_node_count = NUM_NODES
   
     autoscaling {
       max_node_count = MAX_NODES
       location_policy      = "ANY"
     }
     node_config {
       machine_type = MACHINE_TYPE
       reservation_affinity {
         consume_reservation_type = "SPECIFIC_RESERVATION"
         key = "compute.googleapis.com/reservation-name"
         values = [RESERVATION_LABEL_VALUES]
       }
       spot = true
     }
   
     placement_policy {
       type = "COMPACT"
       tpu_topology = TPU_TOPOLOGY
     }
   }
   

   Substitua:

   • NODE_POOL_RESOURCE_NAME: o nome do recurso do pool de nós no modelo do Terraform.
   • PROJECT_ID: o ID do projeto.
   • CLUSTER_NAME: o nome do cluster atual ao qual o pool de nós será adicionado.
   • POOL_NAME: o nome do pool de nós a ser criado.
   • CLUSTER_LOCATION: local do Compute do cluster. Recomendamos ter um cluster regional para aumentar a confiabilidade do plano de controle do Kubernetes. Também é possível usar um cluster zonal. Para saber mais, consulte Selecionar uma versão de TPU e topologia.
   • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.
   • NUM_NODES: o número de nós no pool de nós. Esse valor precisa ser zero ou o produto do número de chips de TPU dividido por quatro. Isso acontece porque, nas frações da TPU de vários hosts, cada nó de TPU tem quatro chips. Por exemplo, se TPU_TOPOLOGY for 4x8, haverá 32 chips, o que significa que NUM_NODES precisa ser 8. Para saber mais sobre topologias de TPU, use a tabela em Mapeamento da configuração de TPU.
   • TPU_TOPOLOGY: indica a topologia física desejada para a fração de TPU. O formato da topologia depende da versão da TPU usada:
     • Para a TPU v4: defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4.
     • Para a TPU v5e: defina a topologia em duas tuplas ({A}x{B}), por exemplo, 2x2.

   Também é possível usar as seguintes variáveis:

   • RESERVATION_NAME: se você usar a reserva de TPU, esta será a lista de rótulos dos recursos de reserva a serem usados ao criar o pool de nós. Para saber como preencher o RESERVATION_LABEL_VALUES no campo reservation_affinity, consulte Provedor do Terraform.
   • autoscaling: adicionar um pool de nós com escalonamento automático ativado. Quando o GKE escalona um pool de nós de uma fração da TPU de vários hosts, ele escalona horizontalmente e de maneira atomizada esse pool de nós de zero até o tamanho máximo.
     • MAX_NODES: o tamanho máximo do pool de nós. Ele precisa ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada VM.
   • spot: permite que o pool de nós use VMs spot para os nós de TPU. Isso não pode ser alterado após a criação do pool de nós. Para mais informações, consulte VMs spot.

Console

Para criar um pool de nós com TPUs:

1. Acesse a página Google Kubernetes Engine no console do Google Cloud.

   Acessar o Google Kubernetes Engine

2. Na lista de clusters, clique no nome do cluster que você quer modificar.

3. Clique em Adicionar pool de nós.

4. Na seção Detalhes do pool de nós, marque a caixa Especificar locais do nó.

5. Selecione a zona com base na versão de TPU que você quer usar:

   • Para a TPU v4, use us-central2-b.
   • Os tipos de máquina TPU v5e que começam com ct5l- nunca são de vários hosts.
   • Para os tipos de máquina TPU v5e que começam com ct5lp-, use us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, ou europe-west4-a.
   • Para os tipos de máquina TPU v5p que começam com ct5p-, use us-east1-d, us-east5-a ou us-east5-c.

6. No painel de navegação, clique em Nós.

7. Na seção Configuração da máquina, selecione TPUs.

8. No menu suspenso Série, selecione uma das seguintes opções:

   • CT4P: para a TPU v4.
   • CT5LP: para a TPU v5e.

9. No menu suspenso Tipo de máquina, selecione o nome da máquina que será usada para os nós. Use a tabela Mapeamento da configuração de TPU para saber como definir o tipo de máquina e a topologia de TPU que criam um pool de nós de TPU de vários hosts.

10. No menu suspenso Topologia da TPU, selecione a topologia física para a fração da TPU.

11. Na caixa de diálogo Alterações necessárias, clique em Fazer alterações.

12. Certifique-se de que o Tipo de disco de inicialização seja Disco permanente padrão ou Disco permanente SSD.

13. Como opção, marque a caixa de seleção Ativar nós em VMs do Spot para usar VMs do Spot nos nós do pool de nós.

14. Clique em Criar.

Estado de provisionamento

Se o GKE não conseguir criar o pool de nós de fração de TPU devido à capacidade insuficiente de TPU, o GKE retornará uma mensagem de erro indicando que não é possível criar os nós de TPU por falta de capacidade.

Se você estiver criando um pool de nós de fração de TPU de host único, a mensagem de erro será semelhante a esta:

2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.

Se você estiver criando um pool de nós de fração de TPU com vários hosts, a mensagem de erro será semelhante a esta:

The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.

Sua solicitação de provisionamento de TPU pode permanecer na fila por muito tempo e no estado "Provisionando" enquanto está na fila.

Quando a capacidade fica disponível, o GKE cria os nós restantes que não foram criados.

Caso você precise de capacidade antes, teste as VMs spot. No entanto, é importante lembrar que as VMs spot consomem uma cota diferente das instâncias sob demanda.

É possível excluir a solicitação de TPU na fila excluindo o pool de nós de fatias de TPU.

Executar a carga de trabalho em nós da TPU

Preparação da carga de trabalho

As cargas de trabalho de TPU têm os requisitos de preparação a seguir.

1. Frameworks como JAX, PyTorch e TensorFlow acessam VMs de TPU usando a biblioteca compartilhada libtpu. libtpu inclui o compilador XLA, o software do ambiente de execução da TPU e o driver da TPU. Cada versão do PyTorch e do JAX requer uma determinada versão de libtpu.so. Para usar TPUs no GKE, use as seguintes versões:

   Tipo de TPU	Versão do libtpu.so
   TPU v5e tpu-v5-lite-podslice tpu-v5-lite-device	Versão recomendada do jax[tpu]: v0.4.9 ou posterior Versão recomendada do torchxla[tpuvm]: v2.1.0 ou mais recente
   TPU v5p	Versão recomendada do jax[tpu]: 0.4.19 ou mais recente. Versão recomendada do torchxla[tpuvm]: sugestão de uso de um build de versão noturno em 23 de outubro de 2023.
   TPU v4 tpu-v4-podslice	jax[tpu]: v0.4.4 ou posterior torchxla[tpuvm]: v2.0.0 ou posterior

2. Defina as variáveis de ambiente a seguir para o contêiner que solicita os recursos da TPU:
   • TPU_WORKER_ID: um número inteiro exclusivo para cada pod. Esse ID denota um código de worker exclusivo na fatia de TPU. Os valores aceitos para esse campo variam de zero ao número de pods menos um.
   • TPU_WORKER_HOSTNAMES: uma lista separada por vírgulas de nomes de host ou endereços IP de VM da TPU que precisam se comunicar entre si na fatia. É necessário que haja um nome do host ou endereço IP para cada VM da TPU na fatia. A lista de endereços IP ou nomes do host é ordenada e zero indexada pelo TPU_WORKER_ID.

   O GKE injeta essas variáveis de ambiente automaticamente usando um webhook mutável quando um job é criado com completionMode: Indexed, subdomain e parallelism > 1 e solicitando as propriedades google.com/tpu. O GKE adiciona um serviço sem comando para que os registros DNS sejam adicionados aos pods de apoio do serviço.

   Ao implantar recursos de vários hosts da TPU com o Kuberay, o GKE fornece um webhook implantável como parte dos modelos do Terraform para executar o Ray no GKE. As instruções para executar o Ray no GKE com TPUs podem ser encontradas no Guia do usuário da TPU. O webhook de modificação injeta essas variáveis de ambiente em clusters do Ray solicitando propriedades google.com/tpu e um seletor de nós cloud.google.com/gke-tpu-topology de vários hosts.

   • No manifesto da carga de trabalho, adicione seletores de nós do Kubernetes para garantir que o GKE programe a carga de trabalho da TPU no tipo de máquina e na topologia de TPU definidos:

       nodeSelector:
         cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
         cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
       

     Substitua:

     • TPU_ACCELERATOR: o nome do acelerador de TPU:
       • Para a TPU v4, use tpu-v4-podslice
       • Para os tipos de máquina TPU v5e que começam com ct5l-, use tpu-v5-lite-device
       • Para os tipos de máquina TPU v5e que começam com ct5lp-, use tpu-v5-lite-podslice
       • Para a TPU v5p, use tpu-v5p-slice
     • TPU_TOPOLOGY: a topologia física da fração da TPU. O formato da topologia depende da versão da TPU, da seguinte maneira:
       • TPU v4: defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4.
       • TPU v5e: defina a topologia em duas tuplas ({A}x{B}), por exemplo, 2x2.
       • TPU v5p: defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4.

Depois de concluir a preparação da carga de trabalho, é possível executar um job que usa TPUs.

As seções a seguir mostram exemplos de como executar um job que realiza computação simples com TPUs.

Exemplo 1: executar uma carga de trabalho que exibe o número de chips de TPU disponíveis em um pool de nós de TPU

A carga de trabalho a seguir retorna o número de chips de TPU em todos os nós em uma fração de TPU de vários hosts. Para criar uma fração de vários hosts, a carga de trabalho tem os seguintes parâmetros:

• Versão do TPU: TPU v4
• Topologia: 2x2x4

Essa seleção de versão e topologia resultam em uma fração de vários hosts.

1. Salve o seguinte manifesto como available-chips-multihost.yaml:

   apiVersion: v1
   kind: Service
   metadata:
     name: headless-svc
   spec:
     clusterIP: None
     selector:
       job-name: tpu-available-chips
   ---
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: tpu-available-chips
   spec:
     backoffLimit: 0
     completions: 4
     parallelism: 4
     completionMode: Indexed
     template:
       spec:
         subdomain: headless-svc
         restartPolicy: Never
         nodeSelector:
           cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
           cloud.google.com/gke-tpu-topology: 2x2x4
         containers:
         - name: tpu-job
           image: python:3.10
           ports:
           - containerPort: 8471 # Default port using which TPU VMs communicate
           - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
           securityContext:
             privileged: true
           command:
           - bash
           - -c
           - |
             pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
             python -c 'import jax; print("TPU cores:", jax.device_count())'
           resources:
             requests:
               cpu: 10
               memory: 500Gi
               google.com/tpu: 4
             limits:
               cpu: 10
               memory: 500Gi
               google.com/tpu: 4
   

2. Implante o manifesto:

   kubectl create -f available-chips-multihost.yaml
   

   O GKE executa uma fatia da TPU v4 com quatro VMs da TPU (fatia da TPU de vários hosts). A fração tem 16 chips interconectados.

3. Verifique se o job criou quatro pods:

   kubectl get pods
   

   O resultado será assim:

   NAME                       READY   STATUS      RESTARTS   AGE
   tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
   tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
   tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
   tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
   

4. Consiga os registros de um dos pods:

   kubectl logs POD_NAME
   

   Substitua POD_NAME pelo nome de um dos pods criados. Por exemplo, tpu-job-podslice-0-5cd8r.

   O resultado será assim:

   TPU cores: 16
   

Exemplo 2: executar uma carga de trabalho que exibe o número de chips de TPU disponíveis na VM de TPU

A carga de trabalho a seguir é um pod estático que exibe o número de chips de TPU anexados a um nó específico. Para criar um nó de host único, a carga de trabalho tem os seguintes parâmetros:

• Versão do TPU: TPU v5e
• Topologia: 2x4

Essa seleção de versão e topologia resultam em uma fração de host único.

1. Salve o seguinte manifesto como available-chips-singlehost.yaml:

   apiVersion: v1
   kind: Pod
   metadata:
     name: tpu-job-jax-v5
   spec:
     restartPolicy: Never
     nodeSelector:
       cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       cloud.google.com/gke-tpu-topology: 2x4
     containers:
     - name: tpu-job
       image: python:3.10
       ports:
       - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
       securityContext:
         privileged: true
       command:
       - bash
       - -c
       - |
         pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         python -c 'import jax; print("Total TPU chips:", jax.device_count())'
       resources:
         requests:
           google.com/tpu: 8
         limits:
           google.com/tpu: 8
   

2. Implante o manifesto:

   kubectl create -f available-chips-singlehost.yaml
   

   O GKE provisiona nós com oito frações de TPU de host único que usam a TPU v5e. Cada VM de TPU tem oito chips (fração de TPU de host único).

3. Consiga os registros do pod:

   kubectl logs tpu-job-jax-v5
   

   O resultado será assim:

   Total TPU chips: 8
   

Fazer upgrade de pools de nós usando aceleradores (GPUs e TPUs)

O GKE faz upgrade automático de clusters Padrão, incluindo pools de nós. Também é possível fazer upgrade manual dos pools de nós se quiser que seus nós usem uma versão mais recente antecipadamente. Para controlar como os upgrades funcionam no cluster, use os canais de lançamento, janelas de manutenção e exclusões e sequenciamento de lançamento.

Também é possível configurar um estratégia de upgrade de nós para o pool de nós, como upgrades súbitos ou upgrades azuis-verdes. Ao configurar essas estratégias, é possível garantir que o upgrade dos pools de nós ocorra de uma maneira que atinja o equilíbrio ideal entre velocidade e interrupção do seu ambiente. Para pools de nós de fração de TPU de vários hosts, em vez de usar a estratégia de upgrade de nós configurada, o GKE recria atomicamente todo o pool de nós em uma única etapa. Para saber mais, consulte a definição de atomicidade em Terminologia relacionada à TPU no GKE.

O uso de uma estratégia de upgrade de nós exigirá temporariamente que o GKE provisione recursos adicionais, dependendo da configuração. Se o Google Cloud tiver capacidade limitada para os recursos do pool de nós, por exemplo, se você estiver enfrentando erros de disponibilidade de recursos ao tentar criar mais nós com GPUs ou TPUs, consulte Fazer upgrade em um ambiente limitado de recursos.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste guia, considere excluir os pools de nós da TPU que não têm mais cargas de trabalho programadas. Se as cargas de trabalho em execução precisarem ser encerradas normalmente, use kubectl drain para limpá-las antes de excluir o nó.

1. Exclua um pool de nós da TPU:

   gcloud container node-pools delete POOL_NAME \
       --location=LOCATION \
       --cluster=CLUSTER_NAME
   

   Substitua:

   • POOL_NAME: o nome do pool de nós.
   • CLUSTER_NAME: o nome do cluster.
   • LOCATION: o local de computação do cluster.

Outras configurações

Nas seções a seguir, descrevemos as configurações adicionais que você pode aplicar às cargas de trabalho da TPU.

Vários setores

É possível agregar frações menores em um Multislice para lidar com cargas de trabalho de treinamento maiores. Para mais informações, consulte TPUs com vários setores no GKE.

Migrar sua reserva da TPU

Se você tiver reservas de TPU atuais, primeiro será necessário migrá-las para um novo sistema de reservas do Compute Engine. Também é possível criar um sistema de reservas com base no Compute Engine, em que nenhuma migração é necessária. Para saber como migrar suas reservas de TPU, consulte Reserva de TPU.

Geração de registros

Os registros emitidos por contêineres em execução em nós do GKE, incluindo VMs de TPU, sãocoletou pelo agente do Logging do GKE, enviado para o Logging e sãovisível no Logging do Google Analytics.

Usar o provisionamento automático de nós do GKE

É possível configurar o GKE para criar e excluir automaticamente pools de nós e atender às demandas de recursos das cargas de trabalho de TPU. Para mais informações, consulte Como configurar Cloud TPUs.

Reparo automático de nó de TPU

Se um nó da TPU em um pool de nós da fatia da TPU de vários hosts não estiver íntegro, todo o pool de nós será recriado. As condições que resultam em nós da TPU não íntegros incluem:

• Qualquer nó da TPU com condições de nó comuns.
• Qualquer nó de TPU com uma TPU não alocável é maior que zero.
• Qualquer instância de VM de TPU interrompida (devido à preempção) ou encerrada.
• Manutenção do nó: se algum nó da TPU (VM) em um pool de nós da fatia da TPU de vários hosts ficar inativo para manutenção do host, o GKE recriará toda a fatia da TPU.

Veja o status do reparo (incluindo o motivo da falha) no histórico de operações. Se a falha for causada por cota insuficiente, entre em contato com o representante da conta do Google Cloud para aumentar a cota correspondente.

Configurar o encerramento do nó da TPU

Visualização

Esse recurso está sujeito aos "Termos de ofertas pré-GA" na seção "Termos de Serviço gerais" dos Termos específicos de serviço. Os recursos pré-GA estão disponíveis "no estado em que se encontram" e podem ter suporte limitado. Para mais informações, consulte as descrições da fase de lançamento.

Nos clusters do GKE com o plano de controle em execução na versão 1.29.1-gke.1425000 ou posterior, os nós da TPU são compatíveis com sinais SIGTERM que alertam o nó sobre um encerramento iminente. A notificação de encerramento iminente pode ser configurada em até cinco minutos em nós da TPU.

É possível configurar o GKE para encerrar as cargas de trabalho de aprendizado de máquina normalmente dentro desse período de notificação. Durante o encerramento sem falhas, as cargas de trabalho podem executar processos de limpeza, como armazenar dados da carga de trabalho para reduzir a perda de dados. Para receber o tempo máximo de notificação, no manifesto do pod, defina o campo spec.terminationGracePeriodSeconds como 300 segundos (cinco minutos) da seguinte maneira:

    spec:
      terminationGracePeriodSeconds: 300

O GKE se esforça para encerrar esses pods da melhor maneira possível e executar a ação de encerramento definida por você, por exemplo, salvando um estado de treinamento.

Observação: pods removidos por encerramento de nó normal são contabilizados nas remoções permitidas configuradas no PDB. Se você definir um PodDisruptionBudget, defina um maxUnavailable alto o suficiente ou um minAvailable baixo o suficiente para permitir que o GKE conclua a remoção do pod dentro da janela de tempo da notificação de encerramento. Caso contrário, talvez o GKE precise violar os parâmetros de disponibilidade do PDB ao remover os pods para que o encerramento continue.

Executar contêineres sem o modo privilegiado

Observação: os contêineres em execução nos nós do GKE 1.28 ou mais recente não precisam ter o modo privilegiado ativado para acessar TPUs.

Se o nó de TPU estiver executando versões anteriores à 1.28, leia a seguinte seção:

Um contêiner em execução em uma VM de TPU precisa de acesso a limites mais altos de memória bloqueada para que o driver possa se comunicar com os chips de TPU por acesso direto à memória (DMA, na sigla em inglês). Para ativar esse recurso, é preciso configurar uma ulimit mais alta. Se você quiser reduzir o escopo de permissão no contêiner, siga estas etapas:

1. Edite o securityContext para incluir os seguintes campos:

   securityContext:
     capabilities:
       add: ["SYS_RESOURCE"]
   

2. Aumente o ulimit executando o seguinte comando dentro do contêiner antes de configurar as cargas de trabalho para usar os recursos da TPU:

   ulimit -l 68719476736
   

Observação: na TPU v5e, a execução de contêineres sem o modo privilegiado está disponível nos clusters nas versões 1.27.4-gke.900 e posteriores.

Observabilidade e métricas

Painel

Na página Clusters do Kubernetes no Console do Google Cloud, a guia Observabilidade exibe as métricas de observabilidade da TPU. Para mais informações, consulte Métricas de observabilidade do GKE.

O painel da TPU será preenchido apenas se você tiver métricas do sistema ativadas no cluster do GKE.

Métricas do ambiente de execução

No GKE 1.27.4-gke.900 ou mais recente, as cargas de trabalho da TPU que usam o JAX 0.4.14 ou mais recente e especificam containerPort: 8431 exportam as métricas de utilização da TPU como métricas do sistema do GKE. As métricas a seguir estão disponíveis no Cloud Monitoring para monitorar o desempenho do ambiente de execução da carga de trabalho da TPU

• Ciclo de trabalho: porcentagem de tempo durante o período de amostragem anterior (60 segundos) em que os TensorCores estavam ativamente em processamento em um chip de TPU. Uma porcentagem maior significa melhor uso da TPU.
• Uso da memória: quantidade de memória do acelerador alocada em bytes. Amostras coletadas a cada 60 segundos.
• Capacidade de memória: memória total do acelerador em bytes. Amostras coletadas a cada 60 segundos.

Essas métricas estão localizadas no esquema de nó (k8s_node) e contêiner de contêiner (k8s_container) do Kubernetes.

Contêiner do Kubernetes:

• kubernetes.io/container/accelerator/duty_cycle
• kubernetes.io/container/accelerator/memory_used
• kubernetes.io/container/accelerator/memory_total

Nó do Kubernetes:

• kubernetes.io/node/accelerator/duty_cycle
• kubernetes.io/node/accelerator/memory_used
• kubernetes.io/node/accelerator/memory_total

Métricas de hospedagem

No GKE 1.28.1-gke.1066000 ou mais recente, a VM da TPU exporta as métricas de utilização da TPU como métricas do sistema do GKE. As seguintes métricas estão disponíveis no Cloud Monitoring para monitorar o desempenho do host da TPU:

• Uso do TensorCore: porcentagem atual do TensorCore que está sendo utilizada. O valor do TensorCore é igual à soma das unidades de multiplicação de matriz (MXUs) mais a unidade vetorial. O valor de utilização do TensorCore é a divisão das operações do TensorCore que foram realizadas no período de amostragem anterior (60 segundos) pelo número de operações do TensorCore com suporte no mesmo período. Um valor maior significa melhor utilização.
• Uso de largura de banda da memória: porcentagem atual da largura de banda da memória do acelerador que está sendo usada. Calculada pela divisão da largura de banda da memória usada durante um período de amostragem (60 segundos) pela largura de banda máxima aceita no mesmo período.

Essas métricas estão localizadas no esquema de nó (k8s_node) e contêiner de contêiner (k8s_container) do Kubernetes.

Contêiner do Kubernetes:

• kubernetes.io/container/accelerator/tensorcore_utilization
• kubernetes.io/container/accelerator/memory_bandwidth_utilization

Nó do Kubernetes:

• kubernetes.io/container/node/tensorcore_utilization
• kubernetes.io/container/node/memory_bandwidth_utilization

Para mais informações, consulte Métricas do Kubernetes e Métricas do sistema do GKE.

Problemas conhecidos

• O escalonador automático de cluster pode calcular incorretamente a capacidade de novos nós da TPU antes que esses nós informem as TPUs disponíveis. Em seguida, o escalonador automático de cluster pode realizar um escalonamento adicional e, como resultado, criar mais nós do que o necessário. O escalonador automático de cluster reduzirá os nós adicionais, se eles não forem necessários, após a operação de redução vertical regular.
• O escalonador automático de cluster cancela o escalonamento vertical dos pools de nós da TPU que permanecem no status de espera por mais de 15 minutos. O autoescalador de cluster tentará realizar essas operações de escalonamento vertical posteriormente. Esse comportamento pode reduzir a capacidade de obtenção de TPU para clientes que não usam reservas.
• As cargas de trabalho que não são da TPU e têm tolerância para o taint da TPU podem impedir a redução em escala do pool de nós se estiverem sendo recriadas durante a diminuição do pool de nós da TPU.

A seguir

• Exibir modelos de linguagem grandes com o Saxml em TPUs
• Saiba como configurar o Ray no GKE com TPUs
• Criar machine learning em larga escala em Cloud TPUs com o GKE
• Exibir modelos de linguagem grandes com o KubeRay em TPUs
• Resolver problemas de TPUs no GKE