Acerca das TPUs no GKE

Esta página descreve como o Cloud TPU funciona com o Google Kubernetes Engine (GKE), incluindo a terminologia, as vantagens das unidades de processamento de tensores (TPUs) e considerações sobre o agendamento de cargas de trabalho. As TPUs são circuitos integrados específicos da aplicação (ASICs) desenvolvidos pela Google para acelerar cargas de trabalho de ML que usam frameworks como o TensorFlow, o PyTorch e o JAX.

Esta página destina-se a administradores e operadores da plataforma, bem como a especialistas em dados e IA que executam modelos de aprendizagem automática (ML) com características como serem de grande escala, de execução prolongada ou dominados por cálculos de matrizes. Para saber mais acerca das funções comuns e das tarefas de exemplo a que fazemos referência no Google Cloud conteúdo, consulte o artigo Funções e tarefas comuns do utilizador do GKE.

Antes de ler esta página, certifique-se de que sabe como funcionam os aceleradores de ML. Para ver detalhes, consulte o artigo Introdução à Cloud TPU.

Vantagens da utilização de TPUs no GKE

O GKE oferece suporte total para a gestão do ciclo de vida dos nós e dos conjuntos de nós da TPU, incluindo a criação, a configuração e a eliminação de VMs da TPU. O GKE também suporta VMs do Spot e a utilização de Cloud TPUs reservadas. Para mais informações, consulte as opções de consumo do Cloud TPU.

Seguem-se algumas vantagens da utilização de TPUs no GKE:

  • Ambiente operacional consistente: pode usar uma única plataforma para todas as cargas de trabalho de aprendizagem automática e outras.
  • Atualizações automáticas: o GKE automatiza as atualizações de versões, o que reduz os custos operacionais.
  • Equilíbrio de carga: o GKE distribui a carga, reduzindo assim a latência e melhorando a fiabilidade.
  • Escalabilidade reativa: o GKE dimensiona automaticamente os recursos de TPU para satisfazer as necessidades das suas cargas de trabalho.
  • Gestão de recursos: com o Kueue, um sistema de filas de tarefas nativo do Kubernetes, pode gerir recursos em vários inquilinos na sua organização através de filas, preemptção, priorização e partilha equitativa.
  • Opções de sandboxing: o GKE Sandbox ajuda a proteger as suas cargas de trabalho com o gVisor. Para mais informações, consulte o artigo GKE Sandbox.

Comece a usar o Ironwood (TPU7x)

O Ironwood (TPU7x) é o TPU de sétima geração da Google, concebido para cargas de trabalho de IA em grande escala. Para mais informações sobre o Ironwood (TPU7x), consulte os seguintes recursos:

Cada uma das seguintes opções de criação de clusters oferece diferentes graus de facilidade e flexibilidade na configuração de clusters 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 validações 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 uma personalização precisa ou uma expansão dos ambientes de produção do GKE existentes

Terminologia relacionada com TPUs no GKE

Esta página usa a seguinte terminologia relacionada com as TPUs:

  • Resiliência da ICI da TPU na nuvem: uma funcionalidade que ajuda a melhorar a tolerância a falhas dos links óticos e dos comutadores de circuitos óticos (OCS) que ligam as TPUs entre cubos. Para mais informações, consulte o artigo Arquitetura da TPU.
  • Cubo de TPU: uma 4x4x4 topologia de chips de TPU interligados. Isto só é aplicável a topologias em tuplos de 3 elementos ({A}x{B}x{C}).
  • Tipo de TPU: o tipo de Cloud TPU, como v5e.
  • Segmento de TPU: uma coleção de chips localizados no mesmo TPU Pod ligados por interconexões de alta velocidade entre chips (ICI). As fatias são descritas em termos de chips ou TensorCores, consoante a versão da TPU.
  • Nó de fatia da TPU: um nó do Kubernetes representado por uma única VM que tem um ou mais chips da TPU interligados.
  • Node pool de fatias de TPUs: um grupo de nós do Kubernetes num cluster que têm todos a mesma configuração de TPUs.
  • Topologia da TPU: o número e a disposição física dos chips da TPU numa fatia da TPU.
  • Atómico: o GKE trata todos os nós interligados como uma única unidade. Durante as operações de escalabilidade, o GKE dimensiona todo o conjunto de nós para 0 e cria novos nós. Se uma máquina no grupo falhar ou terminar, o GKE recria todo o conjunto de nós como uma nova unidade.
  • Imutável: não pode adicionar manualmente novos nós ao conjunto de nós interligados. No entanto, pode criar um novo conjunto de nós com a topologia de TPU pretendida e agendar cargas de trabalho no novo conjunto de nós.

Tipos de pools de nós de fatias de TPU

O GKE suporta dois tipos de conjuntos de nós da TPU:

O tipo e a topologia da TPU determinam se o nó da fatia de TPU pode ser de vários anfitriões ou de um único anfitrião. Recomendamos:

  • Para modelos de grande escala, use nós de fatia de TPU com vários anfitriões.
  • Para modelos de pequena escala, use nós de fatia de TPU de host único.
  • Para a preparação ou a inferência em grande escala, use o Pathways. O Pathways simplifica os cálculos de aprendizagem automática em grande escala, permitindo que um único cliente JAX organize cargas de trabalho em várias grandes fatias de TPUs. Para mais informações, consulte o artigo Pathways.

Pools de nós de fatias de TPU multi-host

Um node pool de fatia de TPU multi-host é um node pool que contém duas ou mais VMs de TPU interligadas. Cada VM tem um dispositivo TPU ligado. As TPUs numa fatia de TPU com vários anfitriões estão ligadas através de uma interligação de alta velocidade (ICI). Depois de criar um conjunto de nós de fatia de TPU com vários anfitriões, não pode adicionar-lhe nós. Por exemplo, não pode criar um conjunto de nós v4-32 e, posteriormente, adicionar um nó do Kubernetes (VM da TPU) ao conjunto de nós. Para adicionar uma fatia de TPU a um cluster do GKE, tem de criar um novo conjunto de nós.

As VMs num conjunto de nós de fatia de TPU com vários anfitriões são tratadas como uma única unidade atómica. Se o GKE não conseguir implementar um nó na fatia, não são implementados nós no nó da fatia de TPU.

Se um nó numa fatia de TPU com vários anfitriões precisar de reparação, o GKE encerra todas as VMs na fatia de TPU, forçando a remoção de todos os pods do Kubernetes na carga de trabalho. Depois de todas as VMs na fatia de TPU estarem em funcionamento, os pods do Kubernetes podem ser agendados nas VMs na nova fatia de TPU.

O diagrama seguinte mostra uma fatia de TPU de vários anfitriões v5litepod-16 (v5e). Esta fatia de TPU tem quatro VMs. Cada VM na fatia de TPU tem quatro chips de TPU v5e ligados com interconexões de alta velocidade (ICI), e cada chip de TPU v5e tem um TensorCore:

Diagrama de fatia de TPU com vários anfitriões

O diagrama seguinte mostra um cluster do GKE que contém uma fatia de TPU v5litepod-16 (v5e) (topologia: 4x4) e uma fatia de TPU v5litepod-8 (v5e) (topologia: 2x4):

Diagrama do pod TPU v5e

Node pools de fatia de TPU de anfitrião único

Um conjunto de nós de divisão de anfitrião único é um conjunto de nós que contém uma ou mais VMs de TPU independentes. Cada VM tem um dispositivo TPU ligado. Embora as VMs num conjunto de nós de fatia de anfitrião único possam comunicar através da rede do centro de dados (DCN), as UTPs anexadas às VMs não estão interligadas.

O diagrama seguinte mostra um exemplo de uma fatia de TPU de host único que contém sete v4-8 máquinas:

Diagrama do conjunto de nós de fatia de anfitrião único

Caraterísticas das TPUs no GKE

As TPUs têm características únicas que requerem um planeamento e uma configuração especiais.

Consumo de TPUs

Para otimizar a utilização e o custo dos recursos, ao mesmo tempo que equilibra o desempenho da carga de trabalho, o GKE suporta as seguintes opções de consumo de TPUs:

Para escolher a opção de consumo que cumpre os requisitos da sua carga de trabalho, consulte o artigo Acerca das opções de consumo de aceleradores para cargas de trabalho de IA/ML no GKE.

Antes de usar as 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 numa fatia de TPU. O GKE aprovisiona uma fatia de TPU em topologias bidimensionais ou tridimensionais, consoante a versão da TPU. Especifica uma topologia como o número de chips de TPU em cada dimensão da seguinte forma:

Para a TPU v4, v5p e Ironwood (TPU7x) (pré-visualização) agendadas em pools de nós de fatias de TPU com vários anfitriões, define a topologia em tuplos de 3 elementos ({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 conjunto de nós. Por exemplo, pode definir pequenas topologias com menos de 64 chips de TPU com formas de topologia, como 2x2x2, 2x2x4 ou 2x4x4. Se usar topologias maiores com mais de 64 chips de TPU, os valores que atribui a {A}, {B} e {C} têm de cumprir as seguintes condições:

  • {A}, {B} e {C} têm de ser múltiplos de quatro.
  • A topologia mais extensa suportada para a v4 é 12x16x16 e para a v5p é 16x16x24.
  • Os valores atribuídos têm de manter o padrão A ≤ B ≤ C. Por exemplo, 4x4x8 ou 8x8x8.

Tipo de máquina

Os tipos de máquinas que suportam 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 fatia de nó, como ct<version>-hightpu-<node-chip-count>t. Por exemplo, o tipo de máquina tpu7x-standard-4t suporta o Ironwood (TPU7x) e contém 4 chips de TPU.

Modo privilegiado

Se usar versões do GKE anteriores à 1.28, tem de configurar os contentores com capacidades especiais para aceder às TPUs. Nos clusters do modo padrão, pode usar o modo privilegiado para conceder este acesso. O modo privilegiado substitui muitas das outras definições de segurança no securityContext. Para ver detalhes, consulte o artigo Execute contentores sem o modo privilegiado.

As versões 1.28 e posteriores não requerem o modo privilegiado nem capacidades especiais.

Como funcionam as TPUs no GKE

A gestão e a priorização de recursos do Kubernetes tratam as VMs em TPUs da mesma forma que outros tipos de VMs. Para pedir chips de TPU, use o nome do recurso google.com/tpu:

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

Quando usar as TPUs no GKE, considere as seguintes características das TPUs:

  • Uma VM pode aceder a até 8 chips de TPU.
  • Uma fatia de TPU contém um número fixo de chips de TPU, sendo que o número depende do tipo de máquina de TPU que escolher.
  • O número de google.com/tpu pedidos tem de ser igual ao número total de chips de TPU disponíveis no nó da fatia de TPU. Qualquer contentor num pod do GKE que peça TPUs tem de consumir todos os chips de TPU no nó. Caso contrário, a implementação falha porque o GKE não consegue 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 fatia de TPU com quatro chips de TPU cada, num total de oito chips de TPU. Com este tipo de máquina, pode:
    • Não é possível implementar um pod do GKE que requer oito chips de TPU nos nós neste conjunto de nós.
    • Pode implementar dois pods que requerem quatro chips de TPU cada, cada pod num dos dois nós neste conjunto 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 esta configuração tem de pedir quatro chips de TPU em cada réplica, para uma a quatro réplicas.
  • Nos clusters padrão, é possível agendar vários pods do Kubernetes numa VM, mas apenas um contentor em cada pod pode aceder aos chips TPU.
  • Para criar pods kube-system, como kube-dns, cada cluster Standard tem de ter, pelo menos, um conjunto de nós de fatia não TPU.
  • Por predefinição, os nós de fatia da TPU têm a google.com/tpu contaminação que impede o agendamento de cargas de trabalho que não sejam da TPU nos nós de fatia da TPU. As cargas de trabalho que não usam TPUs são executadas em nós sem TPUs, libertando recursos de computação em nós de fatias de TPU para código que usa TPUs. Tenha em atenção que a contaminação não garante que os recursos da TPU sejam totalmente utilizados.
  • O GKE recolhe os registos emitidos por contentores executados em nós de segmentos da TPU. Para saber mais, consulte o artigo Registo.
  • As métricas de utilização da TPU, como o desempenho de tempo de execução, estão disponíveis no Cloud Monitoring. Para saber mais, consulte o artigo Observabilidade e métricas.
  • Pode colocar as suas cargas de trabalho de TPU em sandbox com o GKE Sandbox. O GKE Sandbox funciona com modelos de TPU v4 e posteriores. Para saber mais, consulte o artigo GKE Sandbox.

Criação automática de node pools com TPUs

A criação automática de conjuntos de nós suporta os seguintes TPUs na nuvem apenas em versões específicas do GKE:

  • TPU v3: 1.31.0 e posterior.
  • TPU v5 e TPU v4: 1.29.0 e posterior.
  • TPU Trillium: 1.31.1-gke.1146000 e posterior.
  • Ironwood (TPU7x) (pré-visualização): 1.34.1-gke.2541000 ou posterior.

Outros tipos de Cloud TPU são suportados em todas as versões do GKE.

Escala automática do node pool do Cloud TPU

O GKE dimensiona os node pools da TPU do Google Cloud criados automaticamente ou manualmente que usam o dimensionador automático de clusters de uma das seguintes formas:

  • Pool de nós de fatia de TPU de host único: o GKE adiciona ou remove nós da TPU no pool de nós existente. O conjunto de nós pode conter qualquer número de nós de TPU entre zero e o tamanho máximo do conjunto de nós, conforme determinado pelas flags de escalabilidade automática --max-nodes e --total-max-nodes. Todos os nós da TPU no conjunto de nós têm o mesmo tipo de máquina e topologia. Para mais informações sobre como criar um node pool de fatia de TPU de host único, consulte o artigo Crie um node pool de fatia de TPU de host único.
  • Pool de nós de fatia de TPU com vários anfitriões: o GKE aumenta automaticamente o número de nós do pool de nós de zero para o número de nós necessários para satisfazer a topologia da TPU. Por exemplo, com um conjunto de nós da TPU que tenha o tipo de máquina ct5lp-hightpu-4t e uma topologia de 16x16, o conjunto de nós tem sempre 64 nós ou zero nós. O GKE reduz a escala do conjunto de nós se não existirem cargas de trabalho de TPU no conjunto de nós. Para reduzir a escala do node pool, o GKE despeja todos os pods agendados e remove todos os nós no node pool. Para mais informações sobre como criar um node pool de fatia de TPU com vários anfitriões, consulte o artigo Crie um node pool de fatia de TPU com vários anfitriões.

Política de carga de trabalho

A política de carga de trabalho só é suportada para o Ironwood (TPU7x).

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

O Ironwood (TPU7x) suporta a política de carga de trabalho de Alto débito. Esta política permite-lhe fazer o seguinte:

  • Executar cargas de trabalho distribuídas que requerem um elevado débito, como computação de alto desempenho (HPC) ou preparação de aprendizagem automática (ML). A infraestrutura subjacente está configurada para colocar as VMs da TPU no mesmo local, de modo a reduzir a latência da rede entre as VMs.
  • Defina a estratégia de manutenção para VMs de TPU de modo a minimizar as interrupções da carga de trabalho.

Programação da recolha

O agendamento da recolha só é suportado para a TPU Trillium.

Na TPU Trillium, pode usar a programação da recolha para agrupar nós de fatias de TPU. O agrupamento destes nós de fatias de TPU facilita o ajuste do número de réplicas para satisfazer a procura da carga de trabalho.O Google Cloud controla as atualizações de software para garantir que existem sempre fatias suficientes na coleção para publicar tráfego.

O TPU Trillium suporta o agendamento da recolha para pools de nós de anfitrião único e vários anfitriões que executam cargas de trabalho de inferência. A descrição seguinte explica como o comportamento do agendamento da recolha depende do tipo de fatia de TPU que usa:

  • Segmento de TPU com vários anfitriões: o GKE agrupa segmentos de TPU com vários anfitriões para formar uma coleção. Cada GKE node pool é uma réplica nesta coleção. Para definir uma coleção, crie uma fatia de TPU com vários anfitriões e atribua um nome exclusivo à coleção. Para adicionar mais fatias de TPU à coleção, crie outro conjunto de nós de fatias de TPU multi-anfitrião com o mesmo nome da coleção e tipo de carga de trabalho.
  • Segmento de TPU de anfitrião único: o GKE considera o conjunto de nós do segmento de TPU de anfitrião único como uma coleção. Para adicionar mais fatias de TPU à coleção, pode redimensionar o conjunto de nós de fatias de TPU de anfitrião único.

O agendamento da recolha tem as seguintes limitações:

  • Só pode agendar recolhas para a TPU Trillium.
  • Só pode definir coleções durante a criação do conjunto de nós.
  • As VMs Spot não são suportadas.
  • As coleções que contêm pools de nós de fatias de TPUs com vários anfitriões têm de usar o mesmo tipo de máquina, topologia e versão para todos os pools de nós na coleção.

Pode configurar o agendamento da recolha nos seguintes cenários:

O que se segue?

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