Introdução às TPUs no GKE

Os clientes do Google Kubernetes Engine (GKE) agora podem criar pools de nós do Kubernetes com fatias TPU v4 e v5e. Para mais informações sobre TPUs, consulte Arquitetura do sistema.

Ao trabalhar com o GKE, primeiro é necessário criar um cluster do GKE. Em seguida, adicione pools de nós ao cluster. Os pools de nós do GKE são coleções de VMs que compartilham os mesmos atributos. Para cargas de trabalho de TPU, os pools de nós consistem em VMs de TPU.

Tipos de pool de nós

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

• Fração de TPU de vários hosts
• Fração de TPU de host único

Pool de nós de frações 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 TPU conectado a ela. As TPUs em uma fração de vários hosts são conectadas por uma interconexão de alta velocidade (ICI). Depois que um pool de nós de fatia 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, mais tarde, adicionar outro nó do Kubernetes (VM TPU) ao pool de nós. Para adicionar outra fatia de TPU a um cluster do GKE, é necessário criar um novo pool de nós.

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

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

O diagrama a seguir mostra um exemplo de um slice de TPU de vários hosts v5litepod-16 (v5e). Essa fatia tem quatro VMs da TPU. Cada VM de TPU tem quatro chips TPU v5e conectados com interconexões de alta velocidade (ICI), e cada chip TPU v5e tem um TensorCore.

O diagrama a seguir mostra um cluster do GKE com uma fatia de TPU v5litepod-16 (v5e) (topologia: 4x4) e uma fatia de TPU v5litepod-8 (v5e) (topologia: 2x4):

Para conferir um exemplo de execução de uma carga de trabalho em uma fração de TPU de vários hosts, consulte Executar a carga de trabalho em TPUs.

Pools de nós de frações de TPU de host único

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

O diagrama a seguir mostra um exemplo de uma fatia de TPU de host único com sete máquinas v4-8:

Para conferir um exemplo de execução de uma carga de trabalho em uma fração de TPU de host único, consulte Executar a carga de trabalho em TPUs.

Tipos de máquina TPU para pools de nós do GKE

Antes de criar pools de nós, é preciso escolher a versão e o tamanho da fração de TPU necessária para a carga de trabalho. A TPU v4 tem suporte no GKE Standard versão 1.26.1-gke.1500 e mais recentes, v5e no GKE Standard versão 1.27.2-gke.2100 e mais recentes e v5p no GKE Standard versão 1.28.3-gke.1024000 e mais recentes.

A TPU v4, v5e e v5p são compatíveis com o Autopilot do GKE versão 1.29.2-gke.1521000 e mais recentes.

Para mais informações sobre as especificações de hardware das diferentes versões de TPU, consulte Arquitetura do sistema. Ao criar um pool de nós de TPU, selecione um tamanho de fatia de TPU (uma topologia de TPU) com base no tamanho do modelo e na quantidade de memória necessária. O tipo de máquina especificado ao criar os pools de nós depende da versão e do tamanho das fatias.

v5e

Confira a seguir os tipos de máquina e as topologias da TPU v5e com suporte para casos de uso de treinamento e inferência:

Tipo de máquina	topologia	Número de chips do TPU	Número de VMs	Caso de uso recomendado
ct5lp-hightpu-1t	1x1	1	1	Treinamento e inferência de host único
ct5lp-hightpu-4t	2x2	4	1	Treinamento e inferência de host único
ct5lp-hightpu-8t	2x4	8	1	Treinamento e inferência de host único
ct5lp-hightpu-4t	2x4	8	2	Treinamento e inferência de vários hosts
ct5lp-hightpu-4t	4x4	16	4	Treinamento em grande escala, inferência em vários hosts
ct5lp-hightpu-4t	4x8	32	8	Treinamento em grande escala, inferência em vários hosts
ct5lp-hightpu-4t	8x8	64	16	Treinamento em grande escala, inferência em vários hosts
ct5lp-hightpu-4t	8x16	128	32	Treinamento em grande escala, inferência em vários hosts
ct5lp-hightpu-4t	16x16	256	64	Treinamento em grande escala, inferência em vários hosts

O Cloud TPU v5e é um produto de treinamento e inferência combinado. Os jobs de treinamento são otimizados para capacidade e disponibilidade, enquanto os jobs de inferência são otimizados para latência. Para mais informações, consulte Tipos de aceleradores de treinamento v5e e Tipos de aceleradores de inferência v5e.

As máquinas TPU v5e estão disponíveis em us-west4-a, us-east5-b e us-east1-c. Os clusters do GKE Standard precisam executar a versão 1.27.2-gke.2100 ou mais recente do plano de controle. O GKE Autopilot precisa executar a versão 1.29.2-gke.1521000 ou mais recente do plano de controle. Para mais informações sobre a v5e, consulte Treinamento do Cloud TPU v5e.

Comparação entre tipos de máquina:

Tipo de máquina	ct5lp-hightpu-1t	ct5lp-hightpu-4t	ct5lp-hightpu-8t
Número de chips v5e	1	4	8
Número de vCPUs	24	112	224
RAM (GB)	48	192	384
Número de nós NUMA	1	1	2
Probabilidade de preempção	Alta	Média	Baixo

Para criar espaço para VMs com mais chips, o programador do GKE pode interromper e reprogramar VMs com menos chips. Portanto, as VMs de 8 chips têm mais probabilidade de preemptar VMs de 1 e 4 chips.

v4 e v5p

Confira abaixo os tipos de máquina TPU v4 e v5p:

Tipo de máquina	Número de vCPUs	Memória (GB)	Número de nós NUMA
ct4p-hightpu-4t	240	407	2
ct5p-hightpu-4t	208	448	2

Ao criar uma fatia de TPU v4, use o tipo de máquina ct4p-hightpu-4t, que tem um host e contém quatro chips. Consulte Topologias v4 e Arquitetura do sistema de TPU para mais informações. Os tipos de máquinas de fatia TPU v4 estão disponíveis em us-central2-b. Seus clusters do GKE Standard precisam executar a versão 1.26.1-gke.1500 ou mais recente do plano de controle. Os clusters do GKE Autopilot precisam executar a versão 1.29.2-gke.1521000 ou mais recente do plano de controle.

Ao criar uma fatia de TPU v5p, use o tipo de máquina ct5p-hightpu-4t, que tem um host e contém quatro chips. Os tipos de máquina de fatia TPU v5p estão disponíveis em us-west4-a e us-east5-a. Os clusters do GKE Standard precisam executar a versão 1.28.3-gke.1024000 ou mais recente do plano de controle. O Autopilot do GKE precisa ser executado na versão 1.29.2-gke.1521000 ou posterior. Para mais informações sobre o v5p, consulte Introdução ao treinamento v5p.

Limitações e problemas conhecidos

• Número máximo de pods do Kubernetes: é possível executar no máximo 256 pods do Kubernetes em uma única VM de TPU.
• Somente reservas ESPECÍFICAS: ao usar TPUs no GKE, SPECIFIC é o único valor aceito para a flag --reservation-affinity do comando gcloud container node-pools create.
• Apenas a variante de VMs spot de TPUs preemptivas é compatível: as VMs spot são semelhantes às VMs preemptivas e estão sujeitas às mesmas limitações de disponibilidade, mas não têm uma duração máxima de 24 horas.
• Sem suporte à alocação de custos: a alocação de custos do GKE e a medição de uso não incluem dados sobre o uso ou os custos das TPUs.
• O escalonador automático pode calcular a capacidade: o escalonador automático de cluster pode calcular a capacidade incorretamente para novos nós que contêm VMs TPU antes que esses nós estejam disponíveis. Em seguida, o escalonador automático de cluster pode realizar um escalonar verticalmente adicional e, como resultado, criar mais nós do que o necessário. O escalonador automático de cluster vai reduzir escala vertical os nós adicionais, se eles não forem necessários, após a operação de reduzir escala vertical regular.
• O escalonador automático cancela o escalonamento vertical: o escalonador automático de cluster cancela o escalonamento vertical dos pools de nós de TPU que permanecem no status de espera por mais de 10 horas. O autoescalador de cluster vai tentar realizar essas operações de escalonar verticalmente mais tarde. Esse comportamento pode reduzir a capacidade de obtenção de TPU para clientes que não usam reservas.
• O taint pode impedir a redução em escala: as cargas de trabalho que não são da TPU e têm tolerância para o taint da TPU podem impedir a reduzir escala vertical do pool de nós se forem recriadas durante a diminuição do pool de nós da TPU.

Garantir cotas suficientes de TPU e GKE

Talvez seja necessário aumentar algumas cotas relacionadas ao GKE nas regiões em que os recursos são criados.

As cotas a seguir têm valores padrão que provavelmente precisarão ser aumentados:

• Cota SSD de disco permanente (GB): o disco de inicialização de cada nó do Kubernetes requer 100 GB por padrão. Portanto, essa cota precisa ser definida pelo menos tão alta quanto o número máximo de nós do GKE que você prevê que serão criados * 100 GB.
• Cota de endereços IP em uso: cada nó do Kubernetes consome um endereço IP. Portanto, essa cota precisa ser definida pelo menos tão alta quanto o número máximo de nós do GKE que você prevê que serão criados.

Para solicitar um aumento de cota, consulte Solicitar uma cota maior. Para mais informações sobre os tipos de cotas de TPU, consulte Cota de TPU.

Pode levar alguns dias para que seus pedidos de aumento de cota sejam aprovados. Se você tiver dificuldades para aprovar suas solicitações de aumento de cota em alguns dias, entre em contato com a Equipe de Contas do Google.

Migrar sua reserva da TPU

Se você não pretende usar uma reserva de TPU com TPUs no GKE, pule esta seção e acesse Criar um cluster do Google Kubernetes Engine.

Para usar TPUs reservadas com o GKE, primeiro é necessário migrar a reserva de TPU para um novo sistema de reservas do Compute Engine.

Há várias informações importantes sobre essa migração:

• A capacidade da TPU migrada para o novo sistema de reservas do Compute Engine não pode ser usada com a API de recursos em fila do Cloud TPU. Se você pretende usar recursos de fila da TPU com sua reserva, será necessário migrar uma parte dela para o novo sistema de reservas do Compute Engine.
• Nenhuma carga de trabalho pode ser executada ativamente em TPUs quando elas são migradas para o novo sistema de reservas baseado no Compute Engine.
• Selecione um horário para realizar a migração e trabalhe com a equipe da sua conta do Google Cloud para programá-la. A janela de migração precisa ser durante o horário comercial (de segunda a sexta-feira, das 9h às 17h no horário de Brasília).

Criar um cluster do Google Kubernetes Engine

Consulte Criar um cluster na documentação do Google Kubernetes Engine.

Criar um pool de nós de TPU

Consulte Criar um pool de nós na documentação do Google Kubernetes Engine.

Execução sem o modo privilegiado

Se você quiser reduzir o escopo de permissão no contêiner, consulte o Modo de privilégio de TPU.

Executar cargas de trabalho em pools de nós de TPU

Consulte Executar cargas de trabalho do GKE em TPUs na documentação do Google Kubernetes Engine.

seletores de nodes

Para que o Kubernetes programe sua carga de trabalho em nós que contêm VMs TPU, especifique dois seletores para cada carga de trabalho no manifesto do Google Kubernetes Engine:

• Defina cloud.google.com/gke-accelerator-type como tpu-v5-lite-podslice, tpu-v5p-slice ou tpu-v4-podslice.
• Defina cloud.google.com/gke-tpu-topology como a topologia de TPU do nó.

Importante: é obrigatório que os pods do Kubernetes tenham esses dois rótulos especificados no nó nodeSelector ou nodeAffinity do manifesto da carga de trabalho.

As seções Carga de trabalho de treinamento e Carga de trabalho de inferência contêm exemplos de manifestos que ilustram o uso desses seletores de nó.

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. Para mais informações, consulte Considerações sobre a programação de cargas de trabalho na documentação do GKE.

Reparo de nós

Se um nó em um pool de nós de fração de TPU com vários hosts não estiver íntegro, o GKE recria todo o pool de nós. Para mais informações, consulte Reparo automático de nós na documentação do GKE.

Multislice: vá além de uma única fatia

É possível agregar frações menores em um Multislice para lidar com cargas de trabalho de treinamento maiores. Para mais informações, consulte Cloud TPU com vários setores.

Tutoriais de carga de trabalho de treinamento

Estes tutoriais se concentram em treinar cargas de trabalho em uma fração de TPU de vários hosts (por exemplo, quatro máquinas v5e). Eles abrangem os seguintes modelos:

• Modelos FLAX do Hugging Face: treinar a difusão em Pokémon
• PyTorch/XLA: GPT2 no WikiText

Fazer o download de recursos do tutorial

Faça o download dos scripts Python e das especificações YAML do tutorial para cada modelo pré-treinado com o seguinte comando:

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

Criar e se conectar a um cluster

Crie um cluster regional do GKE para que o plano de controle do Kubernetes seja replicado em três zonas, oferecendo maior disponibilidade. Crie o cluster em us-west4, us-east1 ou us-central2, dependendo da versão do TPU que você está usando. Para mais informações sobre TPUs e zonas, consulte Regiões e zonas do Cloud TPU.

O comando a seguir cria um novo cluster regional do GKE inscrito no canal de lançamento rápido com um pool de nós que inicialmente contém um nó por zona. O comando também ativa a Federação de Identidade da Carga de Trabalho para GKE e os recursos do driver CSI do Cloud Storage FUSE no cluster porque as cargas de trabalho de inferência de exemplo neste guia usam buckets do Cloud Storage para armazenar modelos pré-treinados.

gcloud container clusters create cluster-name \
  --region your-region \
  --release-channel rapid \
  --num-nodes=1 \
  --workload-pool=project-id.svc.id.goog \
  --addons GcsFuseCsiDriver

Para ativar a federação de identidade da carga de trabalho para o GKE e os recursos do driver CSI do Cloud Storage FUSE para clusters atuais, execute o seguinte comando:

gcloud container clusters update cluster-name \
  --region your-region \
  --update-addons GcsFuseCsiDriver=ENABLED \
  --workload-pool=project-id.svc.id.goog

As cargas de trabalho de exemplo são configuradas com as seguintes suposições:

• o pool de nós está usando tpu-topology=4x4 com quatro nós
• o pool de nós está usando machine-type ct5lp-hightpu-4t

Execute o comando a seguir para se conectar ao cluster recém-criado:

gcloud container clusters get-credentials cluster-name \
--location=cluster-region

Modelos FLAX do Hugging Face: treine a difusão em Pokémon

Este exemplo treina o modelo de difusão estável do HuggingFace usando o conjunto de dados Pokémon.

O modelo Stable Diffusion é um modelo latente de texto para imagem que gera imagens fotorrealistas a partir de qualquer entrada de texto. Para mais informações sobre a difusão estável, consulte:

• Visão geral do Stable Diffusion
• Fonte de código da Stable Diffusion

Criar imagem do Docker

O Dockerfile está localizado na pasta ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/.

Antes de executar o comando abaixo, verifique se a conta tem as permissões adequadas para que o Docker envie para o repositório.

Crie e envie a imagem do Docker:

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/
docker build -t gcr.io/project-id/diffusion:latest .
docker push gcr.io/project-id/diffusion:latest

Implantar a carga de trabalho

Crie um arquivo com o conteúdo a seguir e nomeie-o como tpu_job_diffusion.yaml. Preencha o campo de imagem com a imagem que você acabou de criar.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-diffusion
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-diffusion
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (e.g. 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-diffusion
        image: gcr.io/${project-id}/diffusion:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU usage metrics, if supported
        command:
        - bash
        - -c
        - |
          cd examples/text_to_image
          python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Em seguida, implante usando:

kubectl apply -f tpu_job_diffusion.yaml

Limpeza

Depois que o job terminar de ser executado, você poderá excluí-lo usando:

kubectl delete -f tpu_job_diffusion.yaml

PyTorch/XLA: GPT2 no WikiText

Este tutorial mostra como executar o GPT2 em TPUs v5e usando o HuggingFace no PyTorch/XLA usando o conjunto de dados wikitext.

Criar imagem do Docker

O Dockerfile está localizado na pasta ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/.

Antes de executar o comando abaixo, verifique se a conta tem as permissões adequadas para que o Docker envie para o repositório.

Crie e envie a imagem do Docker:

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/
docker build -t gcr.io/project-id/gpt:latest .
docker push gcr.io/project-id/gpt:latest

Implantar a carga de trabalho

Copie o YAML abaixo e salve-o em um arquivo chamado tpu_job_gpt.yaml. Preencha o campo de imagem com a imagem que você acabou de criar.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-gpt
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-gpt
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (for example, 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      volumes:
      # Increase size of tmpfs /dev/shm to avoid OOM.
      - name: shm
        emptyDir:
          medium: Memory
          # consider adding `sizeLimit: XGi` depending on needs
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-gpt
        image: gcr.io/$(project-id)/gpt:latest
        ports:
        - containerPort: 8479
        - containerPort: 8478
        - containerPort: 8477
        - containerPort: 8476
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        env:
        - name: PJRT_DEVICE
          value: 'TPU'
        - name: XLA_USE_BF16
          value: '1'
        command:
        - bash
        - -c
        - |
          numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py   --num_cores 4 examples/pytorch/language-modeling/run_clm.py    --num_train_epochs 3 --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 16    --per_device_eval_batch_size 16 --do_train --do_eval  --output_dir /tmp/test-clm     --overwrite_output_dir --config_name my_config_2.json --cache_dir /tmp --tokenizer_name gpt2  --block_size 1024 --optim adafactor --adafactor true --save_strategy no --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json
        volumeMounts:
        - mountPath: /dev/shm
          name: shm
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Implante o fluxo de trabalho usando:

kubectl apply -f tpu_job_gpt.yaml

Limpeza

Depois que o job terminar de ser executado, você poderá excluí-lo usando:

kubectl delete -f tpu_job_gpt.yaml

Tutorial: cargas de trabalho de inferência de host único

Este tutorial mostra como executar uma carga de trabalho de inferência de host único em TPUs do GKE v5e para modelos pré-treinados com JAX, TensorFlow e PyTorch. De modo geral, há quatro etapas separadas a serem realizadas no cluster do GKE:

1. Crie um bucket do Cloud Storage e configure o acesso a ele. Você usa um bucket do Cloud Storage para armazenar o modelo pré-treinado.

2. Faça o download e converta um modelo pré-treinado em um compatível com TPU. Aplique um pod do Kubernetes que faz o download do modelo pré-treinado, usa o conversor da Cloud TPU e armazena os modelos convertidos em um bucket do Cloud Storage usando o driver FUSE CSI do Cloud Storage. O Cloud TPU Converter não requer hardware especializado. Neste tutorial, mostramos como fazer o download do modelo e executar o Cloud TPU Converter no pool de nós de CPU.

3. Inicie o servidor para o modelo convertido. Aplique uma implantação que ofereça o modelo usando uma estrutura de servidor com suporte do volume armazenado no PersistentVolume ReadOnlyMany (ROX). As réplicas de implantação precisam ser executadas em um pool de nós de fatia v5e com um pod do Kubernetes por nó.

4. Implante um balanceador de carga para testar o servidor de modelo. O servidor é exposto a solicitações externas usando o serviço LoadBalancer. Um script Python foi fornecido com um exemplo de solicitação para testar o servidor de modelo.

O diagrama a seguir mostra como as solicitações são roteadas pelo balanceador de carga.

Observação: esses modelos são apenas exemplos de referência e não são otimizados para desempenho. Elas não são recomendadas para uso com serviços de comparação de mercado ou de produção.

Exemplos de implantação do servidor

Esses exemplos de cargas de trabalho são configurados com as seguintes suposições:

• O cluster está em execução com um pool de nós de TPU v5 com três nós
• O pool de nós está usando o tipo de máquina ct5lp-hightpu-1t em que:
  • a topologia é 1x1
  • o número de chips do TPU é 1

O manifesto do GKE a seguir define uma única implantação de servidor host.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: bert-deployment
spec:
  selector:
    matchLabels:
      app: tf-bert-server
  replicas: 3 # number of nodes in node pool
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
      labels:
        app: tf-bert-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1  # target topology
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice  # target version
      containers:
      - name: serve-bert
        image: us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
        env:
        - name: MODEL_NAME
          value: "bert"
        volumeMounts:
        - mountPath: "/models/"
          name: bert-external-storage
        ports:
        - containerPort: 8500
        - containerPort: 8501
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        resources:
          requests:
            google.com/tpu: 1 # TPU chip request
          limits:
            google.com/tpu: 1 # TPU chip request
      volumes:
      - name: bert-external-storage
        persistentVolumeClaim:
          claimName: external-storage-pvc

Se você estiver usando um número diferente de nós no pool de nós do TPU, mude o campo replicas para o número de nós.

Se o cluster padrão executar o GKE versão 1.27 ou anterior, adicione o seguinte campo ao manifesto:

spec:
  securityContext:
    privileged: true

Não é necessário executar pods do Kubernetes no modo privilegiado no GKE versão 1.28 ou mais recente. Para mais detalhes, consulte Executar contêineres sem o modo privilegiado.

Se você estiver usando um tipo de máquina diferente:

• Defina cloud.google.com/gke-tpu-topology como a topologia do tipo de máquina que você está usando.
• Defina os dois campos google.com/tpu em resources para corresponder ao número de chips do tipo de máquina correspondente.

Configuração

Faça o download dos scripts Python e manifestos YAML do tutorial usando o seguinte comando:

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

Acesse o diretório single-host-inference:

cd ai-on-gke/gke-tpu-examples/single-host-inference/

Configurar o ambiente Python

Os scripts do Python usados neste tutorial exigem a versão 3.9 ou mais recente. Instale o requirements.txt para cada tutorial antes de executar os scripts de teste do Python.

Se você não tiver a configuração adequada do Python no seu ambiente local, use o Cloud Shell para fazer o download e executar os scripts do Python neste tutorial.

Configurar o cluster

1. Crie um cluster usando o tipo de máquina e2-standard-4.

   gcloud container clusters create cluster-name \
   --region your-region \
   --release-channel rapid \
   --num-nodes=1 \
   --machine-type=e2-standard-4 \
   --workload-pool=project-id.svc.id.goog \
   --addons GcsFuseCsiDriver

2. Crie o pool de nós de TPU de host único.

Os exemplos de cargas de trabalho pressupõem o seguinte:

• Seu cluster está em execução com um pool de nós TPU v5e com três nós.
• O pool de nós da TPU está usando o tipo de máquina ct5lp-hightpu-1t.

Se você estiver usando uma configuração de cluster diferente da descrita anteriormente, edite o manifesto de implantação do servidor.

Para a demonstração da difusão estável do JAX, você vai precisar de um pool de nós de CPU com um tipo de máquina que tenha 16 Gi ou mais de memória disponível (por exemplo, e2-standard-4). Isso é configurado no comando gcloud container clusters create ou adicionando um pool de nós adicional ao cluster com o seguinte comando:

gcloud beta container node-pools create your-pool-name \
  --zone=your-cluster-zone \
  --cluster=your-cluster-name \
  --machine-type=e2-standard-4 \
  --num-nodes=1

Substitua:

• your-pool-name: o nome do pool de nós a ser criado.
• your-cluster-zone: a zona em que o cluster foi criado.
• your-cluster-name: o nome do cluster em que o pool de nós será adicionado.
• your-machine-type: o tipo de máquina dos nós a serem criados no pool de nós.

Configurar o armazenamento de modelos

Há várias maneiras de armazenar seu modelo para veiculação. Neste tutorial, usaremos a seguinte abordagem:

• Para converter o modelo pré-treinado para funcionar em TPUs, vamos usar uma nuvem privada virtual com suporte do Persistent Disk com acesso ReadWriteMany (RWX).
• Para veicular o modelo em vários TPUs de host único, usaremos a mesma VPC com suporte do bucket do Cloud Storage.

Execute o comando abaixo para criar um bucket do Cloud Storage.

gcloud storage buckets create gs://your-bucket-name \
  --project=your-bucket-project-id \
  --location=your-bucket-location

Substitua:

• your-bucket-name: o nome do bucket do Cloud Storage.
• your-bucket-project-id: o ID do projeto em que você criou o bucket do Cloud Storage.
• your-bucket-location: o local do bucket do Cloud Storage. Para melhorar o desempenho, especifique o local em que o cluster do GKE está em execução.

Siga as etapas abaixo para dar acesso ao bucket ao seu cluster do GKE. Para simplificar a configuração, os exemplos a seguir usam o namespace padrão e a conta de serviço padrão do Kubernetes. Para mais detalhes, consulte Configurar o acesso aos buckets do Cloud Storage usando a Federação de Identidade da Carga de Trabalho para GKE do GKE.

1. Crie uma conta de serviço do IAM para seu aplicativo ou use uma conta de serviço do IAM atual. É possível usar qualquer conta de serviço do IAM no projeto do bucket do Cloud Storage.

   Observação: se você estiver usando uma conta de serviço do IAM, pule esta etapa.

   gcloud iam service-accounts create your-iam-service-acct \
   --project=your-bucket-project-id

   Substitua:

   • your-iam-service-acct: o nome da nova conta de serviço do IAM.
   • your-bucket-project-id: o ID do projeto em que você criou a conta de serviço do IAM. A conta de serviço do IAM precisa estar no mesmo projeto que o bucket do Cloud Storage.

2. Verifique se a conta de serviço do IAM tem os papéis de armazenamento necessários.

   gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
   --member "serviceAccount:[e-mail]" \
   --role "roles/storage.objectAdmin"

   Substitua:

   • your-bucket-name: o nome do bucket do Cloud Storage.
   • your-iam-service-acct: o nome da nova conta de serviço do IAM.
   • your-bucket-project-id: o ID do projeto em que você criou a conta de serviço do IAM.

3. Permita que a conta de serviço do Kubernetes personifique a conta de serviço do IAM. Para isso, adicione uma vinculação de política do IAM entre as duas contas de serviço. Essa vinculação permite que a conta de serviço do Kubernetes atue como a conta de serviço do IAM.

   gcloud iam service-accounts add-iam-policy-binding [e-mail] \
     --role roles/iam.workloadIdentityUser \
     --member "serviceAccount:your-project-id.svc.id.goog[default/default]"

   Substitua:

   • your-iam-service-acct: o nome da nova conta de serviço do IAM.
   • your-bucket-project-id: o ID do projeto em que você criou a conta de serviço do IAM.
   • your-project-id: o ID do projeto em que você criou o cluster do GKE. Os buckets do Cloud Storage e o cluster do GKE podem estar no mesmo projeto ou em projetos diferentes.

4. Anote a conta de serviço do Kubernetes com o endereço de e-mail da conta de serviço do IAM.

   kubectl annotate serviceaccount default \
     --namespace default \
     iam.gke.io/gcp-service-account=[e-mail]

   Substitua:

   • your-iam-service-acct: o nome da nova conta de serviço do IAM.
   • your-bucket-project-id: o ID do projeto em que você criou a conta de serviço do IAM.

5. Execute o comando abaixo para preencher o nome do bucket nos arquivos YAML desta demonstração:

   find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"

   Substitua your-bucket-name pelo nome do bucket do Cloud Storage.

6. Crie o volume permanente e a declaração de volume permanente com o seguinte comando:

   kubectl apply -f pvc-pv.yaml

Inferência e veiculação de modelos JAX

Instale as dependências do Python para executar scripts do Python que enviam solicitações ao serviço de modelo JAX.

pip install -r jax/requirements.txt

Executar a demonstração de veiculação de BERT E2E do JAX:

Esta demonstração usa um modelo BERT pré-treinado do Hugging Face.

Observação: é necessário executar as instruções de configuração antes de executar a demonstração de veiculação de ponta a ponta.

O pod do Kubernetes executa as seguintes etapas:

1. Faz o download e usa o script Python export_bert_model.py dos recursos de exemplo para fazer o download do modelo Bert pré-treinado em um diretório temporário.
2. Usa a imagem do conversor do Cloud TPU para converter o modelo pré-treinado da CPU para TPU e armazena o modelo no bucket do Cloud Storage que você criou durante a configuração.

Esse pod do Kubernetes está configurado para ser executado na CPU padrão do pool de nós. Execute o pod com o seguinte comando:

kubectl apply -f jax/bert/install-bert.yaml

Verifique se o modelo foi instalado corretamente com o seguinte:

kubectl get pods install-bert

Pode levar alguns minutos para que o STATUS leia o Completed.

Iniciar o servidor de modelo do TF para o modelo

Os exemplos de cargas de trabalho neste tutorial pressupõem o seguinte:

• O cluster está em execução com um pool de nós TPU v5 com três nós
• O pool de nós está usando o tipo de máquina ct5lp-hightpu-1t, que contém um chip de TPU.

Se você estiver usando uma configuração de cluster diferente da descrita anteriormente, edite o manifesto de implantação do servidor.

Aplicar a implantação

kubectl apply -f jax/bert/serve-bert.yaml

Verifique se o servidor está em execução com o seguinte:

kubectl get deployment bert-deployment

Pode levar um minuto para que AVAILABLE leia 3.

Aplicar o serviço do balanceador de carga

kubectl apply -f jax/bert/loadbalancer.yaml

Verifique se o balanceador de carga está pronto para tráfego externo com o seguinte:

kubectl get svc tf-bert-service

Pode levar alguns minutos para que o EXTERNAL_IP tenha um IP listado.

Enviar a solicitação ao servidor do modelo

Receba o IP externo do serviço do balanceador de carga:

EXTERNAL_IP=$(kubectl get services tf-bert-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Execute um script para enviar uma solicitação ao servidor:

python3 jax/bert/bert_request.py $EXTERNAL_IP

Saída esperada:

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."

Limpeza

Para limpar os recursos, execute kubectl delete na ordem inversa.

kubectl delete -f jax/bert/loadbalancer.yaml
kubectl delete -f jax/bert/serve-bert.yaml
kubectl delete -f jax/bert/install-bert.yaml

Executar a demonstração de veiculação E2E do JAX Stable Diffusion

Esta demonstração usa o modelo de difusão estável pré-treinado do Hugging Face.

Observação: é necessário executar as instruções de configuração antes de executar a demonstração de veiculação de ponta a ponta.

Exportar o modelo salvo do TF2 compatível com TPU do modelo de difusão estável do Flax

A exportação dos modelos de difusão estáveis exige que o cluster tenha um pool de nós de CPU com um tipo de máquina que tenha 16 GiB ou mais de memória disponível, conforme descrito em Configurar cluster.

O pod do Kubernetes executa as seguintes etapas:

1. Faz o download e usa o script Python export_stable_diffusion_model.py dos exemplos de recursos para fazer o download do modelo de difusão estável pré-treinado em um diretório temporário.
2. Usa a imagem do conversor do Cloud TPU para converter o modelo pré-treinado da CPU para TPU e armazena o modelo no bucket do Cloud Storage que você criou durante a configuração de armazenamento.

Esse pod do Kubernetes está configurado para ser executado no pool de nós de CPU padrão. Execute o pod com o seguinte comando:

kubectl apply -f jax/stable-diffusion/install-stable-diffusion.yaml

Verifique se o modelo foi instalado corretamente com o seguinte:

kubectl get pods install-stable-diffusion

Pode levar alguns minutos para que o STATUS leia o Completed.

Inicie o contêiner do servidor de modelos do TF para o modelo

As cargas de trabalho de exemplo foram configuradas com as seguintes suposições:

• o cluster está em execução com um pool de nós TPU v5 com três nós
• o pool de nós está usando o tipo de máquina ct5lp-hightpu-1t em que:
  • a topologia é 1x1
  • o número de chips do TPU é 1

Se você estiver usando uma configuração de cluster diferente da descrita anteriormente, edite o manifesto de implantação do servidor.

Aplique a implantação:

kubectl apply -f jax/stable-diffusion/serve-stable-diffusion.yaml

Verifique se o servidor está funcionando como esperado:

kubectl get deployment stable-diffusion-deployment

Pode levar um minuto para que AVAILABLE leia 3.

Aplique o serviço do balanceador de carga:

kubectl apply -f jax/stable-diffusion/loadbalancer.yaml

Verifique se o balanceador de carga está pronto para tráfego externo com o seguinte:

kubectl get svc tf-stable-diffusion-service

Pode levar alguns minutos para que o EXTERNAL_IP tenha um IP listado.

Enviar a solicitação ao servidor do modelo

Receba um IP externo do balanceador de carga:

EXTERNAL_IP=$(kubectl get services tf-stable-diffusion-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Executar o script para enviar uma solicitação ao servidor

python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP

Saída esperada:

O comando é Painting of a squirrel skating in New York, e a imagem de saída será salva como stable_diffusion_images.jpg no diretório atual.

Observação: como o modelo pré-treinado é relativamente grande, o serviço de veiculação leva cerca de dois minutos para carregar o modelo. Se você encontrar erros ao enviar uma solicitação para o servidor de modelo, tente novamente após dois minutos.

Limpeza

Para limpar os recursos, execute kubectl delete na ordem inversa.

kubectl delete -f jax/stable-diffusion/loadbalancer.yaml
kubectl delete -f jax/stable-diffusion/serve-stable-diffusion.yaml
kubectl delete -f jax/stable-diffusion/install-stable-diffusion.yaml

Execute a demonstração de exibição do TensorFlow ResNet-50 E2E:

Instale as dependências do Python para executar scripts do Python do tutorial que enviam solicitações para o serviço de modelo do TF.

pip install -r tf/resnet50/requirements.txt

Etapa 1: converter o modelo

Aplicar a conversão de modelo:

kubectl apply -f tf/resnet50/model-conversion.yml

Verifique se o modelo foi instalado corretamente com o seguinte:

kubectl get pods resnet-model-conversion

Pode levar alguns minutos para que o STATUS leia o Completed.

Etapa 2: exibir o modelo com o TensorFlow Serving

Aplique a implantação de disponibilização do modelo:

kubectl apply -f tf/resnet50/deployment.yml

Verifique se o servidor está em execução conforme o esperado com o seguinte comando:

kubectl get deployment resnet-deployment

Pode levar um minuto para que AVAILABLE leia 3.

Aplique o serviço do balanceador de carga:

kubectl apply -f tf/resnet50/loadbalancer.yml

Verifique se o balanceador de carga está pronto para tráfego externo com o seguinte:

kubectl get svc resnet-service

Pode levar alguns minutos para que o EXTERNAL_IP tenha um IP listado.

Etapa 3: enviar a solicitação de teste para o servidor do modelo

Conseguir o IP externo do balanceador de carga:

EXTERNAL_IP=$(kubectl get services resnet-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Execute o script de solicitação de teste (HTTP) para enviar a solicitação ao servidor de modelo.

python3 tf/resnet50/request.py --host $EXTERNAL_IP

Ela será parecida com o exemplo a seguir:

Predict result: ['ImageNet ID: n07753592, Label: banana, Confidence: 0.94921875',
'ImageNet ID: n03532672, Label: hook, Confidence: 0.0223388672', 'ImageNet ID: n07749582,
Label: lemon, Confidence: 0.00512695312

Etapa 4: limpeza

Para limpar os recursos, execute os seguintes comandos kubectl delete:

kubectl delete -f tf/resnet50/loadbalancer.yml
kubectl delete -f tf/resnet50/deployment.yml
kubectl delete -f tf/resnet50/model-conversion.yml

Exclua o pool de nós e o cluster do GKE quando terminar de usá-los.

Inferência e disponibilização de modelos do PyTorch

Instale as dependências do Python para executar scripts do Python que enviam solicitações ao serviço de modelo do PyTorch:

pip install -r pt/densenet161/requirements.txt

Execute a demonstração de veiculação de ponta a ponta do TorchServe Densenet161:

1. Gerar o arquivo do modelo.

   1. Aplique o arquivo do modelo:

   kubectl apply -f pt/densenet161/model-archive.yml

   1. Verifique se o modelo foi instalado corretamente com o seguinte:

   kubectl get pods densenet161-model-archive

   Pode levar alguns minutos para que o STATUS leia o Completed.

2. Disponibilize o modelo com o TorchServe:

   1. Aplique a implantação da disponibilização do modelo:

      kubectl apply -f pt/densenet161/deployment.yml

   2. Verifique se o servidor está em execução conforme o esperado com o seguinte comando:

      kubectl get deployment densenet161-deployment

      Pode levar um minuto para que AVAILABLE leia 3.

   3. Aplique o serviço do balanceador de carga:

      kubectl apply -f pt/densenet161/loadbalancer.yml

      Verifique se o balanceador de carga está pronto para tráfego externo com o seguinte comando:

      kubectl get svc densenet161-service

      Pode levar alguns minutos para que o EXTERNAL_IP tenha um IP listado.

3. Envie uma solicitação de teste para o servidor do modelo:

   1. Receber o IP externo do balanceador de carga:

      EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

   2. Execute o script de solicitação de teste para enviar a solicitação (HTTP) ao servidor de modelo.

      python3 pt/densenet161/request.py --host $EXTERNAL_IP

      Uma resposta como esta será exibida:

      Request successful. Response: {'tabby': 0.47878125309944153, 'lynx': 0.20393909513950348, 'tiger_cat': 0.16572578251361847, 'tiger': 0.061157409101724625, 'Egyptian_cat': 0.04997897148132324

4. Limpe os recursos executando os seguintes comandos kubectl delete:

   kubectl delete -f pt/densenet161/loadbalancer.yml
   kubectl delete -f pt/densenet161/deployment.yml
   kubectl delete -f pt/densenet161/model-archive.yml

   Exclua o pool de nós e o cluster do GKE quando não precisar mais deles.

Como solucionar problemas comuns

Confira informações sobre como resolver problemas do GKE em Resolver problemas de TPU no GKE.

Falha na inicialização da TPU

Se você encontrar o erro abaixo, verifique se está executando o contêiner da TPU no modo privilegiado ou se aumentou o ulimit dentro do contêiner. Para mais informações, consulte Como executar sem o modo privilegiado.

TPU platform initialization failed: FAILED_PRECONDITION: Couldn't mmap: Resource
temporarily unavailable.; Unable to create Node RegisterInterface for node 0,
config: device_path:      "/dev/accel0" mode: KERNEL debug_data_directory: ""
dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true;
could not create driver instance

Programando impasse

Suponha que você tenha dois jobs (Job A e Job B) e que ambos sejam programados em frações de TPU com uma determinada topologia de TPU (por exemplo, v4-32). Suponha também que você tenha duas frações de TPU v4-32 no cluster do GKE. Vamos chamar essas frações de X e Y. Como o cluster tem ampla capacidade de programar os dois jobs, teoricamente eles precisam ser programados rapidamente, um job em cada uma das duas frações de TPU v4-32.

No entanto, sem um planejamento cuidadoso, é possível entrar em um impasse de programação. Suponha que o programador do Kubernetes programe um pod do Kubernetes do job A na fatia X e, em seguida, programe um pod do Kubernetes do job B na fatia X. Nesse caso, dadas as regras de afinidade de pod do Kubernetes para o job A, o programador vai tentar programar todos os pods do Kubernetes restantes para o job A na fatia X. O mesmo para o Job B. Assim, nem o Job A nem o Job B poderão ser programados completamente em uma única fatia. O resultado será um impasse de programação.

Para evitar o risco de um deadlock de programação, use a antiafinidade de pods do Kubernetes com cloud.google.com/gke-nodepool como topologyKey, conforme mostrado neste exemplo:

apiVersion: batch/v1
kind: Job
metadata:
 name: pi
spec:
 parallelism: 2
 template:
   metadata:
     labels:
       job: pi
   spec:
     affinity:
       podAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: In
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
       podAntiAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: NotIn
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
           namespaceSelector:
             matchExpressions:
             - key: kubernetes.io/metadata.name
               operator: NotIn
               values:
               - kube-system
     containers:
     - name: pi
       image: perl:5.34.0
       command: ["sleep",  "60"]
     restartPolicy: Never
 backoffLimit: 4

Como criar recursos de pool de nós de TPU com o Terraform

Também é possível usar o Terraform para gerenciar os recursos do cluster e do pool de nós.

Criar um pool de nós de fração de TPU de vários hosts em um cluster do GKE

Se você tiver um cluster em que quer criar um pool de nós de TPU de vários hosts, use o snippet do Terraform a seguir:

resource "google_container_cluster" "cluster_multi_host" {
  …
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
    workload_pool = "my-gke-project.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "multi_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_multi_host.name
  initial_node_count = 2

  node_config {
    machine_type = "ct4p-hightpu-4t"
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = ["${reservation-name}"]
    }
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = "2x2x2"
  }
}

Substitua os seguintes valores:

• your-project: seu projeto do Google Cloud em que você está executando a carga de trabalho.
• your-node-pool: o nome do pool de nós que você está criando.
• us-central2: a região em que você está executando a carga de trabalho.
• us-central2-b: a zona em que você está executando a carga de trabalho.
• your-reservation-name: o nome da sua reserva.

Criar um pool de nós de fração de TPU de host único em um cluster do GKE

Use o seguinte snippet do Terraform:

resource "google_container_cluster" "cluster_single_host" {
  …
  cluster_autoscaling {
    autoscaling_profile = "OPTIMIZE_UTILIZATION"
  }
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
  workload_pool = "${project-id}.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "single_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_single_host.name
  initial_node_count = 0
  autoscaling {
    total_min_node_count = 2
    total_max_node_count = 22
    location_policy      = "ANY"
  }

  node_config {
    machine_type = "ct4p-hightpu-4t"
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }
}

Substitua os seguintes valores:

• your-project: seu projeto do Google Cloud em que você está executando a carga de trabalho.
• your-node-pool: o nome do pool de nós que você está criando.
• us-central2: a região em que você está executando a carga de trabalho.
• us-central2-b: a zona em que você está executando a carga de trabalho.