Planejar TPUs no GKE

Nesta página, descrevemos como planejar o uso de Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) para reduzir o risco de configuração incorreta da TPU, erros de não disponibilidade ou interrupções por falta de cota.

Antes de usar TPUs no GKE, familiarize-se com as definições e terminologia de TPUs no GKE.

Planejar a configuração da TPU

Para trabalhar com TPUs em clusters do GKE, você precisa planejar a configuração deles. Recomendamos que você siga estas etapas:

  1. Escolha um modo de operação do GKE: execute suas cargas de trabalho em TPUs em um cluster do GKE Autopilot ou Standard. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes.
  2. Selecione a versão da TPU: tipos diferentes de TPU têm recursos distintos, como proporções de preço-desempenho, capacidade de processamento de treinamento e latência de disponibilização. Os tipos de TPU afetam as capacidades de memória e CPU disponíveis.
  3. Valide a disponibilidade de TPU: as TPUs estão disponíveis em regiões específicas do Google Cloud. Para usar um tipo de TPU na carga de trabalho do GKE, o cluster precisa estar em uma região compatível com esse tipo.
  4. Escolha a topologia de TPU: a disposição física das TPUs em um fração de TPU. Escolha uma topologia que corresponda aos requisitos de paralelismo do modelo.

Recomendamos que você use as tabelas de referência nesta página para identificar se os pools de nós são nós de fração de TPU de host único ou de vários hosts.

Escolha um modo de operação do GKE

É possível usar TPUs nos modos de operação disponíveis do GKE para clusters:

  • Modo Autopilot (recomendado): o GKE gerencia a infraestrutura subjacente, como configuração de nós, escalonamento automático, upgrades automáticos, configurações de segurança de referência e a configuração de rede de referência. No Autopilot, você escolhe um tipo e uma topologia de TPU e os especifica no manifesto do Kubernetes. O GKE faz o gerenciamento provisionando nós com TPUs e programando as cargas de trabalho.
  • Modo Standard: você gerencia a infraestrutura subjacente, incluindo a configuração de nós individuais.

Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Escolher um modo de operação do GKE.

Escolher a versão da TPU

As VMs em uma fração de TPU têm as seguintes características técnicas:

Piloto automático

Versão da TPU Tipo de máquina Número de vCPUs Memória (GiB) Número de nós NUMA Máximo de chips de TPU em um nó de fração de TPU
TPU v5p
 tpu-v5p-slice 208 448 2 6.144
TPU v5e
 tpu-v5-lite-podslice 24 a 224 48 a 384 1 256
TPU v5e (somente host único)
 tpu-v5-lite-device 24 a 224 48 a 384 1 a 2 8
TPU v4
 tpu-v4-podslice 240 407 2 4.096

Standard

Versão da TPU Tipo de máquina Número de vCPUs Memória (GiB) Número de nós NUMA Probabilidade de ser interrompido
TPU v5p
 ct5p-hightpu-4t 208 448 2
TPU v5e
 ct5l-hightpu-1t 24 48 1 Maior
TPU v5e
 ct5l-hightpu-4t 112 192 1 Média
TPU v5e
 ct5l-hightpu-8t 224 384 2 Menor
TPU v5e
 ct5lp-hightpu-1t 24 48 1 Maior
TPU v5e
 ct5lp-hightpu-4t 112 192 1 Média
TPU v5e
 ct5lp-hightpu-8t 224 384 1 Baixa
TPU v4
 ct4p-hightpu-4t 240 407 2

Analise as especificações e os preços da TPU na Documentação de preços do Cloud TPU para decidir qual configuração de TPU usar.

Limitações

Considere estas limitações ao escolher a TPU a ser usada:

  • A alocação de custos e a medição de uso do GKE não incluem dados sobre o uso ou os custos da TPU v4 reservada.
  • As TPUs v5p e v5e não são compatíveis com riptide/image streaming em us-east5.
  • O escalonamento automático da TPU v5p é compatível com clusters do GKE com planos de controle que executam pelo menos a versão 1.29.2-gke.1035000 ou 1.28.7-gke.1020000.
  • Para reservas de capacidade, use uma reserva específica.

Validar a disponibilidade da TPU no GKE

As TPUs estão disponíveis em regiões específicas do Google Cloud. Para usar um tipo de TPU no cluster do GKE, o cluster precisa estar em uma região compatível com esse tipo.

Piloto automático

Consulte Regiões e zonas de TPU na documentação do Cloud TPU.

Standard

A tabela a seguir lista a disponibilidade de TPU para cada versão de TPU e tipo de máquina:

Versão da TPU Tipo de máquina começando com Versão mínima do GKE Disponibilidade Zona
TPU v5e ct5l- 1.27.2-gke.2100 Disponibilidade geral europe-west4-b
us-central1-a
TPU v5e ct5lp- 1.27.2-gke.2100 Disponibilidade geral europe-west4-a1
us-central1-a1
us-east1-c
us-east5-b1
us-west1-c
us-west4-a
us-west4-b1
TPU v5p ct5p- 1.28.3-gke.1024000 Disponibilidade geral us-east1-d
us-east5-a
us-east5-c
TPU v4 ct4p- 1.26.1-gke.1500 Disponibilidade geral us-central2-b

  1. É possível criar um pool de nós da TPU v5e de host único com um tipo de máquina que começa com ct5lp-, mas não começa com ct5l- em determinadas zonas (europe-west4-a, us-east5-b e us-west4-b). Você pode usar ct5lp-hightpu-4t com uma topologia de pelo menos 2x4 ou maior nessas zonas. Para criar uma TPU v5e de host único na região us-west4, escolha a zona us-west4-a e use tipos de máquina que começam com ct5lp-, como ct5lp-hightpu-1t. Para criar uma TPU v5e de host único nas outras regiões listadas na tabela anterior, use tipos de máquina que começam com ct5l- (como ct5l-hightpu-1t, ct5l-hightpu-4t ou ct5l-hightpu-8t). Observação: os tipos de máquina que começam com ct5l- exigem uma cota diferente em relação aos tipos de máquina que começam com ct5lp-. ↩

Escolher uma topologia

Depois de escolher uma versão da TPU, selecione uma topologia compatível com o tipo de TPU. Dependendo do tipo de TPU, a topologia é bidimensional ou tridimensional. Os requisitos de paralelismo do modelo ajudarão você a decidir sobre uma topologia. É possível identificar o número de chips de TPU na fração ao calcular o produto de cada tamanho na topologia. Por exemplo:

  • 2x2x2 é uma fração de TPU v4 com vários hosts com oito chips
  • 2x2 é uma fração da TPU v5e de host único com quatro chips

Se uma topologia específica for compatível com nós de frações de TPU de host único e de vários hosts, o número de chips de TPU solicitados pela sua carga de trabalho determinará o tipo de host.

Por exemplo, a TPU v5e (tpu-v5-lite-podslice) será compatível com a topologia 2x4 como host único e de vários hosts caso você:

  • Solicite quatro ícones na sua carga de trabalho para receber um nó de vários hosts com quatro chips de TPU.
  • Solicite oito ícones na sua carga de trabalho e receba um host único com 8 chips de TPU.

Use a tabela a seguir para escolher o tipo de máquina e a topologia de TPU para seu caso de uso:

  • Para treinamento ou inferência de modelos em pequena escala, use a TPU v4 ou a TPU v5e com pools de nós de frações de TPU de host único.

  • Para treinamento ou inferência de modelos em grande escala, use a TPU v4 ou a TPU v5e com pools de nós de frações de TPU de vários hosts.

Piloto automático

Versão da TPU Tipo de máquina Topologia Número de chips de TPU em uma fração Número de nós Tipo de pool de nós
TPU v5p tpu-v5p-slice 2x2x1 4 1 Host único
2x2x2 8 2 Vários hosts
2x2x4 16 4 Vários hosts
2x4x4 32 8 Vários hosts
4x4x4 64 16 Vários hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Vários hosts
TPU v5e tpu-v5-lite-podslice2 1x1 1 1 Host único
2x2 4 1
2x4 8 1
2x4 2 1 Vários hosts
4x4 16 4
4x8 32 8
8x8 64 16
8x16 128 32
16x16 256 64
TPU v5e (somente host único) tpu-v5-lite-device 1x1 1 1 Host único
2x2 4 1
2x4 8 1
TPU v4 tpu-v4-podslice2 2x2x1 4 1 Host único
2x2x2 8 2 Vários hosts
2x2x4 16 4 Vários hosts
2x4x4 32 8 Vários hosts
4x4x4 64 16 Vários hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Vários hosts

  1. Calculado pelo produto de topologia dividido por quatro. ↩

    Topologias personalizadas para mais de 64 chips são aceitas. Aplicam-se as seguintes condições:

    • Para mais de 64 chips, {A}, {B} e {C} precisam ser múltiplos de 4
    • A maior topologia é 16x16x24
    • Os valores precisam ser {A}≤{B}≤{C}, como 8x12x16.

  2. Não há suporte para topologias personalizadas.

Depois de escolher um tipo de TPU e uma topologia, especifique-os no manifesto da carga de trabalho. Para instruções, consulte Implantar cargas de trabalho de TPU no Autopilot do GKE.

Standard

Versão da TPU Tipo de máquina Topologia Número de chips do TPU Número de VMs Tipo de pool de nós
TPU v5p ct5p-hightpu-4t 2x2x1 4 1 Host único
2x2x2 8 2 Vários hosts
2x2x4 16 4 Vários hosts
2x4x4 32 8 Vários hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Vários hosts
TPU v5e ct5l-hightpu-1t 1x1 1 1 Host único
ct5l-hightpu-4t 2x2 4 1 Host único
ct5l-hightpu-8t 2x4 8 1 Host único
ct5lp-hightpu-1t 1x1 1 1 Host único
ct5lp-hightpu-4t 2x2 4 1 Host único
ct5lp-hightpu-8t 2x4 8 1 Host único
ct5lp-hightpu-4t 2x4 8 2 Vários hosts
4x4 16 4 Vários hosts
4x8 32 8 Vários hosts
8x8 64 16 Vários hosts
8x16 128 32 Vários hosts
16x16 256 64 Vários hosts
TPU v4 ct4p-hightpu-4t 2x2x1 4 1 Host único
2x2x2 8 2 Vários hosts
2x2x4 16 4 Vários hosts
2x4x4 32 8 Vários hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Vários hosts

  1. Calculado pelo produto de topologia dividido por quatro. ↩

Configurações avançadas

As seções a seguir descrevem as práticas recomendadas de programação para configurações avançadas de TPU.

Reserva da TPU

As reservas de TPU estão disponíveis ao comprar um compromisso. Qualquer reserva de TPU pode ser usada com o GKE.

Ao criar um pool de nós de fração de TPU, use as sinalizações --reservation e --reservation-affinity=specific para consumir uma instância de TPU reservada.

Como fazer o escalonamento automático de TPUs no GKE

O GKE é compatível com Unidades de Processamento de Tensor (TPUs) para acelerar as cargas de trabalho de machine learning. O pool de nós de fração de TPU de host único e o pool de nós de frações de TPU de vários hosts são compatíveis com escalonamento automático e provisionamento automático.

Com a sinalização --enable-autoprovisioning em um cluster do GKE, o GKE cria ou exclui pools de nós de frações de TPU de host único ou de vários hosts com uma versão e topologia de TPU que atende aos requisitos das cargas de trabalho pendentes.

Quando você usa --enable-autoscaling, o GKE escalona o pool de nós com base no tipo, da seguinte maneira:

  • Pool de nós da fração de TPU de host único: o GKE adiciona ou remove nós da TPU no pool de nós atual. O pool de nós pode conter qualquer número de nós da TPU entre zero e o tamanho máximo dele, conforme determinado por --max-nodes e os --total-max-nodes. Quando o pool de nós é escalonado, todos os nós da TPU nele têm o mesmo tipo de máquina e topologia. Para saber mais sobre como criar um pool de nós de fatia de TPU de host único, consulte Criar um pool de nós.

  • Pool de nós de fração de TPU de vários hosts: o GKE escalona atomicamente o pool de nós de zero para o número de nós necessários para satisfazer a topologia de TPU. Por exemplo, com um pool de nós de TPU com um tipo de máquina ct5lp-hightpu-4t e uma topologia de 16x16, o pool de nós contém 64 nós. O escalonador automático do GKE garante que esse pool de nós tenha exatamente 0 ou 64 nós. Ao reduzir o escalonamento, o GKE remove todos os pods programados e drena todo o pool de nós para zero. Para saber mais como criar um pool de nós de fração de TPU de vários hosts, consulte Criar um pool de nós.

CPU para clusters padrão

Esta seção não se aplica aos clusters do Autopilot porque o GKE coloca cada fração de TPU no próprio nó. Para saber mais, consulte Como as TPUs funcionam no modo Autopilot.

Para clusters padrão, considere as seguintes práticas recomendadas de programação.

Para programar uma carga de trabalho que não seja de TPU em uma VM em um nó de fração de TPU, verifique se o pod do GKE pode tolerar o taint google.com/tpu. Se você quiser que a carga de trabalho seja implantada em nós específicos, use seletores de nó.

O gerenciamento de recursos e a prioridade do Kubernetes tratam as VMs nas TPUs da mesma forma que os outros de VM. Para dar prioridade de agendamento a Pods que exigem TPUs em relação a outros Pods nos mesmos nós, solicite a CPU ou memória máxima para essas fatias de TPU. Frações de TPU de baixa prioridade devem fazer o seguinte:

  1. Defina solicitações baixas de CPU e memória para garantir que o nó tenha recursos alocáveis suficientes para as cargas de trabalho da TPU. Para saber mais, consulte Como o Kubernetes aplica solicitações e limites de recursos.
  2. Não há limite de CPU (ilimitado) para garantir que os pods possam ter um burst para usar todos os ciclos não utilizados.
  3. Defina limites de memória adequados para garantir que os pods funcionem corretamente sem o risco de remoção da pressão do nó.

Se um pod do Kubernetes não solicitar CPU e memória (mesmo que esteja solicitando TPUs), o Kubernetes o considerará um pod de melhor esforço e não haverá garantia de que ele precisará de CPU e memória. Somente os pods que solicitam explicitamente a CPU e a memória têm essas garantias. Para uma programação específica do Kubernetes, configure as necessidades do pod com solicitação explícita de CPU e memória. Para mais informações, consulte Gerenciamento de recursos para pods e contêineres.

Para saber mais sobre as práticas recomendadas, consulte Práticas recomendadas do Kubernetes: solicitações e limites de recursos.

Reduzir a interrupção da carga de trabalho

Se você estiver usando TPUs para treinar um modelo de machine learning e a carga de trabalho for interrompida, todo o trabalho realizado desde o último checkpoint será perdido. Para diminuir a probabilidade de interrupção da carga de trabalho, faça o seguinte:

  • Defina uma prioridade mais alta para este job do que todos os outros jobs: se os recursos forem escassos, o programador do GKE antecipa os jobs de menor prioridade para programar um job de prioridade mais alta. Isso também garante que a carga de trabalho com prioridade mais alta receba todos os recursos necessários (até o total de recursos disponíveis no cluster). Para saber mais, consulte Prioridade e preempção do pod.
  • Configurar a exclusão de manutenção: uma exclusão de manutenção é uma janela de tempo não recorrente durante a qual a manutenção automática é proibida. Para saber mais, consulte Exclusões de manutenção.
  • Use pods de tempo de execução prolongado no Autopilot: use pods de tempo de execução estendido por um período de carência de até sete dias antes que o GKE encerre seus pods para redução de escala ou upgrade de nós.

Lidar com interrupções devido à manutenção do nó

Os nós do GKE que hospedam as TPUs estão sujeitos a eventos de manutenção ou outras interrupções que podem causar o encerramento do nó. Nos clusters do GKE com o plano de controle executando a versão 1.29.1-gke.1425000 e posterior, é possível reduzir as interrupções nas cargas de trabalho configurando o GKE para encerrar as cargas de trabalho normalmente.

Para entender, configurar e monitorar eventos de interrupção que podem ocorrer em nós do GKE que executam cargas de trabalho de IA/ML, consulte Gerenciar a interrupção de nós do GKE para GPUs e TPUs.

Maximize a utilização da TPU

Para maximizar seu investimento em TPUs, programe uma combinação de prioridades de job e coloque-as em fila para maximizar o tempo de operação das TPUs. Para programação e preempção no nível do job, você precisa usar um complemento do Kubernetes que orquestra jobs em filas. Recomendamos usar o Kueue para esse caso de uso.

A seguir