Sobre TPUs no GKE

Nesta página, apresentamos o Cloud TPU e ajudamos você a planejar sua configuração do Cloud TPU com o Google Kubernetes Engine (GKE), incluindo a reserva de instâncias de TPU, escalonamento automático, limitações da TPU e considerações sobre programação de carga de trabalho.

As Unidades de Processamento de Tensor (TPUs) são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning (ML) que usam frameworks como oTensorFlow, PyTorch e JAX.

Antes de usar TPUs no GKE, recomendamos que você conclua o seguinte programa de aprendizado:

1. Saiba como os aceleradores de machine learning funcionam com a Introdução à Cloud TPU.
2. Saiba mais sobre a disponibilidade atual da versão da TPU com a arquitetura do sistema do Cloud TPU.

Para saber como configurar a Cloud TPU no GKE, consulte os recursos a seguir:

• Implante cargas de trabalho de TPU no Autopilot do GKE
• Implantar cargas de trabalho de TPU no GKE Standard

Benefícios do uso de TPUs no GKE

O GKE oferece suporte completo para o gerenciamento do ciclo de vida de nós da TPU e do pool de nós, incluindo criação, configuração e exclusão de VMs de TPU. O GKE também oferece suporte a VMs do Spot e usa a Cloud TPU reservada. Os benefícios do uso de TPUs no GKE incluem:

• Ambiente operacional consistente: é possível usar uma única plataforma para todo o machine learning e outras cargas de trabalho.
• Upgrades automáticos: o GKE automatiza as atualizações de versão, o que reduz a sobrecarga operacional.
• Balanceamento de carga: o GKE distribui a latência, reduzindo a carga e melhorando a confiabilidade.
• Escalonamento responsivo: o GKE escalona automaticamente os recursos da TPU para atender às necessidades das cargas de trabalho.
• Gerenciamento de recursos: comKueue, um sistema de enfileiramento de jobs nativo do Kubernetes, é possível gerenciar recursos em vários locatários na sua organização usando enfileiramento, preempção, priorização e compartilhamento justo.

Terminologia relacionada à TPU no GKE

Nesta página, usamos a seguinte terminologia relacionada às TPUs:

• Tipo de TPU: o tipo de Cloud TPU, como v5e.
• Nó de fração da TPU: um nó do Kubernetes representado por uma única VM que tem um ou mais chips da TPU interconectados.
• Pool de nós de fração de TPU: um grupo de nós do Kubernetes em um cluster que todos têm a mesma configuração de TPU.
• Topologia de TPU: o número e a disposição física dos chips de TPU em uma fração de TPU.
• Pool de nós da TPU de host único: um pool de nós que contém um ou mais nós de fração de TPU independentes. As TPUs anexadas a nós diferentes não são interconectadas.
• Pool de nós de fração de TPU de vários hosts: um pool de nós que contém dois ou mais nós de fração de TPU interconectados. Os chips de TPU em um nó de fração de TPU de vários hosts são conectados por interconexões de alta velocidade. Os pools de nós de fatias TPU de vários hosts têm as seguintes características:
  • Atômico: o GKE trata todos os nós interconectados como uma única unidade. Durante as operações de escalonamento, o GKE escalona todo o conjunto de nós para 0 e cria novos nós. Se uma máquina no grupo falhar ou for encerrada, o GKE recria todo o conjunto de nós como uma nova unidade.
  • Imutável: não é possível adicionar manualmente novos nós ao conjunto de nós interconectados. No entanto, é possível criar um novo pool de nós com a topologia de TPU desejada e programar cargas de trabalho no novo pool de nós.

Como funcionam as TPUs no GKE

O gerenciamento de recursos e a prioridade do Kubernetes tratam as VMs nas TPUs da mesma forma que os outros de VM. Você solicita chips de TPU por meio do nome do recurso google.com/tpu:

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

Ao usar TPUs no GKE, considere as seguintes características da TPU:

• Uma VM pode acessar até oito chips de TPU.
• Uma fração de TPU contém um número fixo de chips de TPU, com o número dependendo do tipo de máquina TPU escolhido.
• O número solicitado de google.com/tpu precisa ser igual ao número total de chips TPU disponíveis no nó de fração da TPU. Qualquer contêiner em um pod do GKE que solicite TPUs precisa consumir todos os chips TPU no nó. Caso contrário, a implantação falhará, porque o GKE não pode consumir parcialmente os recursos de TPU. Por exemplo, veja os seguintes cenários:
  • O tipo de máquina ct5l-hightpu-8t tem um único nó de fração de TPU com oito chips de TPU. Assim, em um nó, você:
    • Pode implantar um pod do GKE que requer oito chips de TPU.
    • Não é possível implantar dois pods do GKE que exigem quatro chips de TPU cada.
  • O tipo de máquina ct5lp-hightpu-4t com uma topologia 2x4 contém dois nós de fração de TPU com quatro chips de TPU cada, um total de oito chips de TPU. Com esse tipo de máquina:
    • Não é possível implantar um pod do GKE que requer oito chips do TPU nos nós deste pool.
    • Você pode implantar dois pods que exigem quatro chips de TPU cada, cada um em um dos dois nós neste pool de nós.
  • A TPU v5e com topologia 4x4 tem 16 chips de TPU em quatro nós. A carga de trabalho do GKE Autopilot que seleciona essa configuração precisa solicitar quatro chips de TPU em cada réplica, para de um a quatro réplicas.
• Em clusters padrão, vários pods do Kubernetes podem ser programados em uma VM, mas apenas um contêiner em cada pod pode acessar os ícones da TPU.
• Para criar pods do kube-system, como o kube-dns, cada cluster padrão precisa ter pelo menos um pool de nós de fração não TPU.
• Por padrão, os nós de fração de TPU têm o google.com/tpu taint, que impede que frações de não TPU sejam programadas nos nós de fração de TPU. Cargas de trabalho que não usam TPUs são executadas em nós não TPU, liberando computação em nós de fatias TPU para código que usa TPUs. O taint não garante que os recursos da TPU sejam totalmente utilizados.
• O GKE coleta os registros emitidos por contêineres em execução em nós de fração da TPU. Para saber mais, consulte Logging.
• As métricas de inicialização da TPU, como o desempenho do ambiente de execução, estão disponíveis no Cloud Monitoring. Para saber mais, consulte Observabilidade e métricas.

Planejar a configuração da TPU

Para trabalhar com TPUs em clusters do GKE, é preciso decidir os seguintes parâmetros:

• Tipo de TPU: o tipo de máquina, como ct5l-hightpu-8t. Diferentes tipos de TPU têm diferentes recursos, proporções de preço-desempenho, capacidade de processamento de treinamento e latência de disponibilização. Os tipos de TPU afetam a CPU e as capacidades de memória disponíveis.
• Topologia: a disposição física das TPUs em uma fração de TPU. Cada tipo de TPU aceita uma topologia de TPU 2D ou 3D. Selecione uma topologia que corresponda aos requisitos de paralelismo do modelo.

• Interconexão com TPU: se os nós têm interconexões de alta velocidade. O tipo e a topologia da TPU determinam se é possível ter nós de fração de TPU com vários hosts, que são TPUs em vários nós com interconexões de alta velocidade. Recomendamos o seguinte:

  • Para modelos em grande escala, use nós de fração de TPU de vários hosts
  • Para modelos de pequena escala, use nós de fração de TPU de host único

• Modo privilegiado: substitui muitas das outras configurações de segurança no securityContext. Para acessar TPUs, os contêineres em execução nos nós do GKE em:

  • A versão 1.28 e anteriores precisam ativar o modo privilegiado.
  • As versões 1.28 ou mais recentes não precisam do modo privilegiado.

Escolher uma configuração de TPU para o modo Autopilot do GKE

No modo Autopilot, você escolhe um tipo de TPU e uma topologia e os especifica no manifesto do Kubernetes. O GKE gerencia nós de provisionamento com TPUs e programa suas cargas de trabalho.

Disponibilidade da TPU no Autopilot do GKE

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. Para detalhes, consulte Regiões e zonas de TPU na documentação do Cloud TPU.

Escolher um tipo de TPU no Autopilot

Confira as especificações e os preços do chip de TPU na Documentação de preços do Cloud TPU para ajudar a decidir qual tipo de TPU usar.

Escolher uma topologia para o Autopilot

Depois de escolher um tipo de TPU, selecione uma topologia compatível com ele. 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 oferecer suporte a nós de frações de TPU de host único e de vários hosts, o número de chips TPU solicitados pela sua carga de trabalho determinará o tipo de host que você receberá.

Por exemplo, a TPU v5e (tpu-v5-lite-podslice) é compatível com a topologia 2x4 como host único e de vários hosts. Veja o que acontecerá nestes casos:

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

A tabela a seguir lista cada tipo de TPU, as topologias compatíveis e anotações de uso. Para cada uma dessas topologias, a tabela lista o número de chips de TPU, o número de nós e o tipo de host:

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.

Escolher uma configuração de TPU para o modo GKE Standard

As seções a seguir descrevem as características da TPU a serem consideradas ao planejar e configurar suas cargas de trabalho de TPU no GKE. Para detalhes sobre versões disponíveis, tipos de máquina, topologias válidas e o número de chips de TPU, consulte Mapeamento de configurações de TPU nesta página.

Disponibilidade da TPU no modo GKE Standard

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

As seções a seguir descrevem as características da TPU a serem consideradas ao planejar e configurar suas cargas de trabalho de TPU no GKE. Para detalhes sobre versões disponíveis, tipos de máquina, topologias válidas e o número de chips de TPU, consulte a seção Mapeamento de configurações de TPU nesta página.

Tipo de máquina

Os tipos de máquinas que oferecem suporte aos recursos de TPU seguem uma convenção de nomenclatura que inclui a versão da TPU e o número de chips de TPU por nó, como ct<version>-hightpu-<node-chip-count>t. Por exemplo, o tipo de máquina ct5lp-hightpu-1t é compatível com TPU v5e e contém apenas um chip de TPU.

Topologia

A topologia define a disposição física das TPUs dentro de uma fração da TPU. O GKE provisiona uma fração da TPU em topologias bidimensionais ou tridimensionais, dependendo da versão da TPU. Especifique uma topologia como o número de chips de TPU em cada dimensão:

• Para a TPU v4 e v5p programada em pools de nós de fração da TPU de vários hosts, defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4. O produto de {A}x{B}x{C} define o número de chips de TPU no pool de nós. Por exemplo, é possível definir topologias pequenas menores que 64 chips de TPU com formulários de topologia, como 2x2x2, 2x2x4 ou 2x4x4. Se você usar topologias maiores que 64 chips de TPU, os valores atribuídos a {A}, {B} e {C} precisarão atender às seguintes condições:

  • {A}, {B} e {C} são múltiplos de quatro.
  • A maior topologia compatível com v4 é 12x16x16 e v5p é 16x16x24.
  • Os valores atribuídos mantêm o padrão A ≤ B ≤ C. Por exemplo, 4x4x8 ou 8x8x8.

Mapeamento da configuração 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.

Características da TPU v5e

As máquinas TPU v5e têm as seguintes características técnicas:

Características da TPU v4 e v5p

As máquinas TPU v4p e v5p têm as seguintes características técnicas:

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.

Limitações

Use estas considerações ao planejar o uso de TPUs na sua plataforma:

• Para reservas de capacidade, é preciso usar uma reserva específica.
• 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 1.29.2-gke.1035000 ou 1.28.7-gke.1020000.

Considerações sobre a programação da carga de trabalho

As TPUs têm características únicas que exigem programação e gerenciamento especiais de cargas de trabalho no Kubernetes. Nas seções a seguir, descrevemos as práticas recomendadas de programação.

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 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. Se você quiser a programação e a preempção no nível do job, será necessário usar um complemento para o Kubernetes que organize os jobs em filas. Recomendamos usar o Kueue para esse caso de uso.

A seguir

• Consulte Implantar cargas de trabalho da TPU no GKE para configurar a Cloud TPU com o GKE.
• Conheça as práticas recomendadas para usar o Cloud TPU em tarefas de machine learning.
• Criar machine learning em larga escala em Cloud TPUs com o GKE.
• Exibir modelos de linguagem grandes com o KubeRay em TPUs.