Sobre TPUs no GKE

Este documento descreve como o Cloud TPU funciona com o Google Kubernetes Engine (GKE), incluindo terminologia, os benefícios das Unidades de Processamento de Tensor (TPUs) e considerações sobre programação de carga de trabalho. As 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 ML que usam frameworks como TensorFlow, PyTorch e JAX.

Este documento é destinado a administradores e operadores de plataforma e especialistas em dados e IA que executam modelos de machine learning (ML) com características como grande escala, longa duração ou dominados por cálculos de matriz. Para saber mais sobre papéis comuns e tarefas de exemplo referenciados no Google Cloud conteúdo, consulte Papéis e tarefas de usuário comuns do GKE.

Antes de ler este documento, familiarize-se com o funcionamento dos aceleradores de ML. Para mais detalhes, consulte Introdução ao Cloud TPU.

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. Para mais informações, consulte Opções de consumo do Cloud TPU.

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: Com Kueue, 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.
  • Opções de sandbox:o GKE Sandbox ajuda a proteger suas cargas de trabalho com o gVisor. Para mais informações, consulte GKE Sandbox.

Começar a usar o Ironwood (TPU7x)

O Ironwood (TPU7x) é a TPU de sétima geração do Google, projetada para cargas de trabalho de IA em grande escala. Para mais informações sobre os benefícios do Ironwood (TPU7x), consulte Sobre o Ironwood (TPU7x) no G1.

Terminologia relacionada a TPUs no GKE

Neste documento, usamos a seguinte terminologia relacionada às TPUs:

  • Resiliência da ICI do Cloud TPU:um recurso que ajuda a melhorar a tolerância a falhas de links ópticos e interruptores de circuito óptico (OCS) que conectam as TPUs entre os cubos. Para mais informações, consulte Arquitetura da TPU.
  • Cubo de TPU:uma topologia 4x4x4 de chips de TPU interconectados. Isso só é aplicável a topologias em três tuplas ({A}x{B}x{C}).
  • Tipo de TPU:o tipo de Cloud TPU, como v5e.
  • Fração de TPU:uma coleção de chips no mesmo Pod de TPU conectados por interconexões entre chips (ICI) de alta velocidade. As frações são descritas como chips ou TensorCores, dependendo da versão da TPU.
  • 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.
  • 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.

Tipos de pools de nós de fração de TPU

O GKE oferece suporte a dois tipos de pools de nós de TPU:

O tipo e a topologia da TPU determinam se o nó de fração da TPU pode ser de host único ou de vários hosts. 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.
  • Para treinamento ou inferência em grande escala, use Pathways. O Pathways simplifica os cálculos de machine learning em grande escala, permitindo que um único JAX orquestre cargas de trabalho em várias frações de TPU grandes. Para mais informações, consulte Pathways.

Pools de nós de fração de TPU de vários hosts

Um pool de nós de fração de TPU de vários hosts é um pool de nós que contém duas ou mais VMs de TPU interconectadas. Cada VM tem um dispositivo de TPU conectado a ela. As TPUs em uma fração de TPU de vários hosts são conectadas por uma interconexão de alta velocidade (ICI). Depois que um pool de nós de fração de TPU de vários hosts é criado, não é possível adicionar nós a ele. Por exemplo, não é possível criar um v4-32 pool de nós e adicionar um nó do Kubernetes (VM de TPU) ao pool de nós mais tarde. Para adicionar uma fração de TPU a um cluster do GKE, é necessário criar um novo pool de nós.

As VMs em um pool de nós de fração de TPU de vários hosts são tratadas como uma única unidade atômica. Se o GKE não conseguir implantar um nó na fração, nenhum nó na fração de TPU será implantado.

Se um nó em uma fração de TPU de vários hosts precisar de reparo, o GKE vai desligar todas as VMs na fração de TPU, forçando a remoção de todos os pods do Kubernetes na carga de trabalho. Depois que todas as VMs na fração de TPU estiverem em funcionamento, os pods do Kubernetes poderão ser programados nas VMs na nova fração de TPU.

O diagrama a seguir mostra uma fração de TPU de vários hosts v5litepod-16 (v5e). Essa fração de TPU tem quatro VMs. Cada VM na fração de TPU tem quatro chips de TPU v5e conectados com interconexões de alta velocidade (ICI), e cada chip de TPU v5e tem um TensorCore:

Diagrama de fração de TPU de vários hosts

O diagrama a seguir mostra um cluster do GKE que contém uma fração de TPU v5litepod-16 (v5e) (topologia: 4x4) e uma fração de TPU v5litepod-8 (v5e) (topologia: 2x4):

Diagrama do pod de TPU v5e

Pools de nós de fração de TPU de host único

Um pool de nós de fração de host único é um pool de nós que contém uma ou mais VMs de TPU independentes. Cada VM tem um dispositivo de TPU conectado a ela. Embora as VMs em um pool de nós de fração de host único possam se comunicar pela rede do data center (DCN), as TPUs anexadas às VMs não são interconectadas.

O diagrama a seguir mostra um exemplo de uma fração de TPU de host único que contém sete máquinas v4-8:

Diagrama do pool de nós de fatia de host único

Características das TPUs no GKE

As TPUs têm características únicas que exigem planejamento e configuração especiais.

Consumo de TPU

Para otimizar a utilização de recursos e o custo, equilibrando o desempenho da carga de trabalho, o GKE oferece suporte às seguintes opções de consumo de TPU:

  • Início flexível:para provisionar VMs de início flexível por até sete dias, com o GKE alocando automaticamente o hardware com base na disponibilidade. Para mais informações, consulte Sobre o provisionamento de GPUs, TPUs e H4Ds com o modo de provisionamento de início flexível.
  • VMs do Spot:para provisionar VMs do Spot, é possível receber descontos significativos, mas as VMs do Spot podem ser interrompidas a qualquer momento, com um aviso de 30 segundos. Para mais informações, consulte VMs spot.
  • Reserva adiantada por até 90 dias (no modo de calendário) : para provisionar recursos de TPU por até 90 dias, para um período especificado. Para mais informações, consulte Solicitar TPUs com reserva adiantada no modo de calendário.
  • Reservas de TPU: para solicitar uma reserva adiantada por um ano ou mais.
  • Sob demanda:para consumir TPUs sem organizar a capacidade com antecedência. Antes de solicitar recursos, é necessário ter cota sob demanda suficiente para o tipo e a quantidade específicos de VMs de TPU. A opção sob demanda é a mais flexível. No entanto, não há garantia de que recursos sob demanda suficientes estarão disponíveis para atender à sua solicitação.

A opção sob demanda é o modelo de consumo padrão para TPUs no GKE se você não especificar outra opção. Para escolher a opção de consumo que atende aos requisitos da carga de trabalho, consulte Sobre as opções de consumo de aceleradores para cargas de trabalho de IA/ML no GKE.

Antes de usar TPUs no GKE, escolha a opção de consumo que melhor se adapta aos requisitos da carga de trabalho.

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 da seguinte maneira:

Para a TPU v4, v5p e Ironwood (TPU7x) 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 com menos de 64 chips de TPU com formulários de topologia, como 2x2x2, 2x2x4 ou 2x4x4. Se você usar topologias maiores com mais de 64 chips de TPU, os valores atribuídos a {A}, {B} e {C} precisarão atender às seguintes condições:

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

Nomenclatura do tipo de máquina

A nomenclatura do tipo de máquina para TPUs no GKE varia dependendo do modo de cluster e da versão da TPU:

  • GKE Standard:selecione um tipo de máquina específico do Compute Engine, por exemplo, ct6e-standard-1t para TPU Trillium (v6e).

  • GKE Autopilot:não é possível selecionar tipos de máquina diretamente. Em vez disso, solicite TPUs usando um tipo de acelerador no manifesto da carga de trabalho. Por exemplo, use tpu-v6e-slice para TPU Trillium (v6e) ou tpu-v5-lite-podslice para TPU v5e. O GKE Autopilot provisiona os nós subjacentes com os tipos de máquina apropriados para atender à solicitação.

Para conferir os tipos de máquina exatos disponíveis para cada versão da TPU, consulte as tabelas em Planejar TPUs no GKE.

Modo privilegiado

Se você usar versões do GKE anteriores à 1.28, será necessário configurar seus contêineres com recursos especiais para acessar TPUs. Em clusters do modo Standard, é possível usar o modo privilegiado para conceder esse acesso. O modo privilegiado substitui muitas das outras configurações de segurança no securityContext. Para mais detalhes, consulte Executar contêineres sem modo privilegiado.

As versões 1.28 e mais recentes não exigem modo privilegiado ou recursos especiais.

Como funcionam as TPUs no GKE

O gerenciamento de recursos e a priorização do Kubernetes tratam as VMs nas TPUs da mesma forma que os outros tipos de VM. Para solicitar chips de TPU, use o 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. Considere os seguintes cenários:
    • 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, você:
    • Não é possível implantar um pod do GKE que requer oito chips do TPU nos nós deste pool de nós.
    • É possível implantar dois pods que exigem quatro chips de TPU cada, cada pod 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 chips de 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 cargas de trabalho 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.
  • É possível usar o GKE Sandbox para cargas de trabalho de TPU. O GKE Sandbox funciona com modelos de TPU v4 e mais recentes. Para saber mais, consulte GKE Sandbox.

Criação automática de pool de nós com TPUs

A criação automática de pool de nós oferece suporte apenas às seguintes Cloud TPUs em versões específicas do GKE:

  • TPU v3: 1.31.0 e mais recentes.
  • TPU v5 e TPU v4: 1.29.0 e mais recentes.
  • TPU Trillium: 1.32.0 e mais recentes.
  • Ironwood (TPU7x): 1.34.1-gke.2541000 ou mais recentes.

Outros tipos de Cloud TPU são aceitos em todas as versões do GKE. Para saber mais sobre as versões do GKE disponíveis para TPUs, consulte Validar a disponibilidade da TPU no GKE.

Escalonamento automático do pool de nós do Cloud TPU

O GKE escalona pools de nós do Cloud TPU criados automaticamente ou manualmente que usam o escalonador automático de cluster de uma das seguintes maneiras:

  • 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 the --max-nodes e os --total-max-nodes flags de escalonamento automático. Todos os nós da TPU no pool de nós têm o mesmo tipo de máquina e topologia. Para mais informações sobre como criar um pool de nós de fração de TPU de host único, consulte Criar um pool de nós de fração de TPU de host único.
  • 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 que tem o tipo de máquina ct5lp-hightpu-4t e uma topologia de 16x16, o pool de nós sempre tem 64 nós ou zero nós. O GKE reduz o pool de nós se não houver cargas de trabalho de TPU no pool de nós. Para reduzir o pool de nós, o GKE remove todos os pods programados e todos os nós no pool de nós. Para mais informações sobre como criar um pool de nós de fração de TPU de vários hosts, consulte Criar um pool de nós de fração de TPU de vários hosts.

Programação de coleção

A programação de coleção só é aceita para TPU Trillium.

Na TPU Trillium, é possível usar a programação de coleção para agrupar nós de fração de TPU. O agrupamento desses nós de fração de TPU facilita o ajuste do número de réplicas para atender à demanda da carga de trabalho. Google Cloud controla as atualizações de software para garantir que frações suficientes na coleção estejam sempre disponíveis para atender ao tráfego.

A TPU Trillium oferece suporte à programação de coleção para pools de nós de host único e de vários hosts que executam cargas de trabalho de inferência. A seguir, descrevemos como o comportamento de programação de coleção depende do tipo de fração de TPU usada:

  • Fração de TPU de vários hosts:o GKE agrupa frações de TPU de vários hosts para formar uma coleção. Cada pool de nós do GKE é uma réplica nessa coleção. Para definir uma coleção, crie uma fração de TPU de vários hosts e atribua um nome exclusivo à coleção. Para adicionar mais frações de TPU à coleção, crie outro pool de nós de fração de TPU de vários hosts com o mesmo nome de coleção e tipo de carga de trabalho.
  • Fração de TPU de host único:o GKE considera todo o pool de nós de fração de TPU de host único como uma coleção. Para adicionar mais frações de TPU à coleção, é possível redimensionar o pool de nós de fração de TPU de host único.

A programação de coleção tem as seguintes limitações:

  • Só é possível programar coleções para TPU Trillium.
  • É possível definir coleções apenas durante a criação do pool de nós.
  • As VMs spot não são aceitas.
  • As coleções que contêm pools de nós de fração de TPU de vários hosts precisam usar o mesmo tipo de máquina, topologia e versão para todos os pools de nós na coleção.

É possível configurar a programação de coleção nos seguintes cenários:

A seguir

Para saber como configurar a Cloud TPU no GKE, consulte as seguintes páginas: