Sobre TPUs no GKE

Nesta página, descrevemos 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) desenvolvidos especialmente pelo Google para acelerar cargas de trabalho de ML que usam frameworks como TensorFlow, PyTorch e JAX.

Esta página é destinada 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 conteúdo do Google Cloud, consulte Tarefas e papéis de usuário comuns do GKE.

Antes de ler esta página, confira se você sabe como os aceleradores de ML funcionam. 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 carga, reduzindo a latência 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 o 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 isolamento em 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)

A 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 o Ironwood (TPU7x), consulte os seguintes recursos:

Cada uma das opções a seguir para criação de cluster oferece diferentes graus de facilidade e flexibilidade na configuração do cluster e no agendamento de cargas de trabalho:

  • Use o Accelerated Processing Kit (XPK) para criar rapidamente clusters do GKE e executar cargas de trabalho para provas de conceito e testes. Para mais informações, consulte a documentação do XPK.
  • Use a Google Cloud CLI para criar manualmente a instância do cluster do GKE para personalização ou expansão precisas de ambientes de produção do GKE atuais.

Terminologia relacionada a TPUs no GKE

Nesta página, 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 dos links ópticos e dos interruptores de circuito óptico (OCS) que conectam as TPUs entre os cubos. Para mais informações, consulte Arquitetura de 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:um conjunto de chips localizados 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 nele.

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

O GKE é compatível com dois tipos de pools de nós da TPU:

O tipo e a topologia da TPU determinam se o nó de fração de TPU pode ser de vários ou de um único host. 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 o Pathways. O Pathways simplifica computações de machine learning em grande escala ao permitir que um único cliente JAX orquestre cargas de trabalho em várias frações grandes de TPU. 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 pool de nós v4-32 e depois adicionar um nó do Kubernetes (VM de TPU) a ele. Para adicionar uma fração de TPU a um cluster do GKE, crie 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ó no nó da fração de TPU será implantado.

Se um nó em uma fração de TPU de vários hosts precisar de reparo, o GKE 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 execução, 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 fatia da TPU tem quatro VMs. Cada VM na fração de TPU tem quatro chips TPU v5e conectados com interconexões de alta velocidade (ICI), e cada chip 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 da 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:

Para escolher a opção de consumo que atende aos requisitos da sua 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 sua 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) (pré-lançamento) programadas 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 que têm 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.

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 fração de nó, como ct<version>-hightpu-<node-chip-count>t. Por exemplo, o tipo de máquina tpu7x-standard-4t é compatível com Ironwood (TPU7x) e contém quatro chips de TPU.

Modo privilegiado

Se você usar versões do GKE anteriores à 1.28, configure seus contêineres com recursos especiais para acessar TPUs. Em clusters no modo padrão, é 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 nem 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 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:

  • Uma VM pode acessar até oito chips de TPU.
  • Uma fração de TPU contém um número fixo de chips de TPU, que depende do tipo de máquina de TPU escolhido.
  • O número solicitado de google.com/tpu precisa ser igual ao número total de chips de 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 vai 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ê pode:
    • Não é possível implantar um pod do GKE que requer oito chips de TPU nos nós deste pool.
    • É 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 uma 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 fração de 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 colocar as cargas de trabalho de TPU em sandbox com o GKE Sandbox. 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 é compatível com as seguintes Cloud TPUs apenas em versões específicas do GKE:

  • TPU v3: 1.31.0 e versões mais recentes.
  • TPU v5 e TPU v4: 1.29.0 e versões mais recentes.
  • TPU Trillium: 1.31.1-gke.1146000 e mais recente.
  • Ironwood (TPU7x) (pré-lançamento): 1.34.1-gke.2541000 ou mais recente.

Outros tipos de Cloud TPU são compatíveis com todas as versões do GKE.

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

O GKE escalona verticalmente pools de nós da 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 pelas flags de escalonamento automático --max-nodes e --total-max-nodes. 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 a escala do pool de nós se não houver cargas de trabalho de TPU nele. Para reduzir a escala do pool de nós, o GKE remove todos os pods programados e todos os nós no pool. 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.

Política de carga de trabalho

A política de carga de trabalho só é compatível com o Ironwood (TPU7x).

Uma política de carga de trabalho é um tipo de política de recursos que permite configurar o posicionamento físico da infraestrutura subjacente. Uma política de carga de trabalho oferece opções de configuração hierárquica para um controle mais refinado.

O Ironwood (TPU7x) é compatível com a política de carga de trabalho de alta capacidade de processamento. Com essa política, você pode fazer o seguinte:

  • Execute cargas de trabalho distribuídas que exigem alta capacidade de transferência, como computação de alto desempenho (HPC) ou treinamento de machine learning (ML). A infraestrutura subjacente é configurada para alocar as VMs de TPU no mesmo local e reduzir a latência de rede entre elas.
  • Defina a estratégia de manutenção para VMs de TPU e minimize as interrupções nas cargas de trabalho.

Programação de coleta

O agendamento de coleta só é compatível com a TPU Trillium.

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

A TPU Trillium oferece suporte à programação de coleta para pools de nós de host único e vários hosts que executam cargas de trabalho de inferência. A seguir, descrevemos como o comportamento do agendamento de coleta 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 a ela. Para adicionar mais frações de TPU à coleta, crie outro pool de nós de fração de TPU de vários hosts com o mesmo nome de coleta 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, redimensione o pool de nós de fração de TPU de host único.

O agendamento de coleta tem as seguintes limitações:

  • Só é possível programar coletas para o TPU Trillium.
  • É possível definir coleções apenas durante a criação do pool de nós.
  • As VMs spot não são compatíveis.
  • 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 o agendamento da coleta nos seguintes cenários:

A seguir

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