Esta página foi traduzida pela API Cloud Translation.

Introdução às TPUs no GKE

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

Ao trabalhar com o GKE, primeiro você precisa criar Cluster do GKE. Depois, adicione pools de nós ao aglomerado. Os pools de nós do GKE são coleções de VMs que compartilham os mesmos atributos. Para cargas de trabalho da TPU, os pools de nós são compostos por VMs da TPU.

Tipos de pool de nós

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

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

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

Um pool de nós de fração de TPU com vários hosts é um pool de nós que contém dois ou mais interconectadas da TPU. Cada VM tem um dispositivo TPU conectado a ela. As TPUs do uma fração de vários hosts é conectada por uma interconexão de alta velocidade (ICI). Uma vez que um pool de nós de fração 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 outro (VM da TPU) para o pool de nós. Para adicionar uma fração de TPU adicional a uma cluster do GKE, você precisa 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 fração, não nós na fração serão implantados.

Se um nó em uma fração de TPU de vários hosts precisar ser reparado, o GKE encerrará todas as VMs TPU na fração, forçando todos os pods do Kubernetes na carga de trabalho seja despejada. Assim que todas as VMs de TPU na fração estiverem em execução, a Os pods do Kubernetes podem ser programados nas VMs da TPU na nova fração.

O diagrama a seguir mostra um exemplo de TPU de vários hosts v5litepod-16 (v5e) fatia Essa fração tem quatro VMs de TPU. Cada VM da TPU tem quatro chips da 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 contendo uma Fração v5litepod-16 (v5e) da TPU (topologia: 4x4) e uma v5litepod-8 da TPU (v5e) fatia (topologia: 2x4):

Diagrama do pod da TPU v5e

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

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 um ou mais e VMs de TPU independentes. Cada VM tem um dispositivo TPU conectado a ela. Enquanto as VMs em um pool de nós de fração de host único podem se comunicar pelo de rede (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 fração de TPU de host único com sete Máquinas v4-8:

Diagrama do pool de nós de fração de host único

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

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

Antes de criar pools de nós, é preciso escolher a versão da TPU e o tamanho do Fração de TPU que sua carga de trabalho exige. A TPU v4 é compatível com o GKE Versão padrão 1.26.1-gke.1500 e mais recentes, v5e no GKE Versão padrão 1.27.2-gke.2100 e mais recentes, e v5p na 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 dos diferentes tipos de TPU, versões, consulte Arquitetura do sistema. Quando ao criar um pool de nós de TPU, selecione um tamanho de fração de TPU (uma topologia de TPU) com base na o tamanho do modelo e a quantidade de memória necessária. O tipo de máquina que você especificar ao criar pools de nós depende da versão e do tamanho das frações.

v5e

A seguir estão os tipos de máquina e 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, inferência de host único
`ct5lp-hightpu-4t`	2x2	4	1	Treinamento, inferência de host único
`ct5lp-hightpu-8t`	2x4	8	1	Treinamento, inferência de host único
`ct5lp-hightpu-4t`	2x4	8	2	Treinamento, inferência de vários hosts
`ct5lp-hightpu-4t`	4x4	16	4	Treinamento em grande escala, inferência de vários hosts
`ct5lp-hightpu-4t`	4x8	32	8	Treinamento em grande escala, inferência de vários hosts
`ct5lp-hightpu-4t`	8x8	64	16	Treinamento em grande escala, inferência de vários hosts
`ct5lp-hightpu-4t`	8x16	128	32	Treinamento em grande escala, inferência de vários hosts
`ct5lp-hightpu-4t`	16x16	256	64	Treinamento em grande escala, inferência de vários hosts

O Cloud TPU v5e é um produto combinado de treinamento e inferência. Os jobs de treinamento são otimizados para capacidade e disponibilidade, enquanto os jobs de inferência são otimizados para latência de rede. Para mais informações, consulte Tipos de aceleradores de treinamento v5e e os tipos de acelerador 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 um plano de controle versão 1.27.2-gke.2100 ou posterior. Autopilot do GKE precisa executar o plano de controle na versão 1.29.2-gke.1521000 ou posterior. Para mais informações sobre a v5e, consulte Treinamento da 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 liberar espaço para VMs com mais ícones, o programador do GKE pode reprogramar e reprogramar VMs com menos ícones. As VMs de oito chips têm mais chances de ou de 1 e 4 chips.

v4 e v5p

Veja a seguir 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 fração da TPU v4, use o tipo de máquina ct4p-hightpu-4t, que 1 host e contém 4 ícones. Consulte Topologias v4. e a arquitetura do sistema de TPU informações imprecisas ou inadequadas. Os tipos de máquinas frações de TPU v4 estão disponíveis em us-central2-b. Seu Os clusters do GKE Standard precisam executar um plano de controle versão 1.26.1-gke.1500 ou mais recente. Autopilot do GKE os clusters precisam executar a versão 1.29.2-gke.1521000 ou posterior do plano de controle.

Ao criar uma fração de TPU v5p, use o tipo de máquina ct5p-hightpu-4t, que 1 host e contém 4 ícones. Os tipos de máquina frações de TPU v5p estão disponíveis em us-west4-a e us-east5-a. GKE Standard os clusters precisam executar a versão 1.28.3-gke.1024000 ou posterior do plano de controle. O Autopilot do GKE precisa executar 1.29.2-gke.1521000 ou mais tarde. Para mais informações sobre a v5p, consulte Introdução ao treinamento da v5p (em inglês).

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 da TPU.
Somente reservas ESPECÍFICAS: ao usar TPUs no GKE, SPECIFIC é o único valor compatível com a sinalização --reservation-affinity. do comando gcloud container node-pools create.
Apenas a variante de VMs spot das TPUs preemptivas é compatível: VMs spot são semelhantes às VMs preemptivas e estão sujeitas à mesma disponibilidade com limitações, mas não têm uma duração máxima de 24 horas.
Sem suporte para alocação de custos: alocação de custos do GKE e a medição de uso não inclua 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 de maneira incorreta para novos nós com VMs de TPU antes que esses nós sejam disponíveis. O escalonador automático de clusters pode então realizar escalonar verticalmente verticais adicionais criar mais nós do que o necessário. O escalonador automático de clusters será reduzir escala vertical nós adicionais, se não forem necessários, após a operação normal de reduzir escala vertical vertical.
O escalonador automático cancela o escalonamento vertical: o escalonador automático de cluster cancela o escalonamento vertical da TPU. pools de nós que permanecem no status de espera por mais de 10 horas. cluster O escalonador automático tentará executar essas operações de escalonar verticalmente novamente mais tarde. Esse comportamento pode reduzir a capacidade de obtenção de TPU para clientes que não usam reservas.
Taint pode impedir a redução da escala vertical: cargas de trabalho não TPU que têm uma tolerância para o taint da TPU pode impedir reduzir escala vertical da escala vertical do pool de nós se eles forem recriados durante a drenagem do pool de nós da TPU.

Garanta cotas suficientes de TPU e GKE

Talvez seja necessário aumentar algumas cotas relacionadas ao GKE na as regiões onde os recursos são criados.

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

Cota de SSD do disco permanente (GB): o disco de inicialização de cada nó do Kubernetes requer 100 GB por padrão. Portanto, essa cota deve ser definida pelo menos alto como (o número máximo de nós do GKE que você prevê criação) * 100 GB.
Cota de endereços IP em uso: cada nó do Kubernetes consome um endereço IP. Portanto, essa cota deve ser definida pelo menos no mesmo nível do número máximo nós do GKE que você pretende criar.

Para solicitar aumento da 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 suas solicitações de aumento de cota sejam aprovadas. Se você enfrentar dificuldades para conseguir a aprovação de suas solicitações de aumento de cota em um entre em contato com a Equipe de Contas do Google.

Migrar sua reserva da TPU

Se você não planeja usar uma reserva de TPU atual 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 você precisa migrar sua reserva de TPU para uma nova reserva baseada no Compute Engine. sistema.

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

A capacidade de TPU foi migrada para a nova reserva baseada no Compute Engine sistema não pode ser usado com a API de recursos em fila da Cloud TPU. Se você pretende usar recursos na fila da TPU com sua reserva, precisará migrar uma parte da sua reserva de TPU para o novo Sistema de reservas baseado no Compute Engine.
Nenhuma carga de trabalho pode ser executada ativamente em TPUs quando elas são migradas para o o novo sistema de reservas baseado no Compute Engine.
Selecione um horário para realizar a migração e trabalhar com sua conta do Google Cloud equipe de conta do Google para agendar a migração. A janela de tempo da migração precisa ser durante o horário comercial (de segunda a sexta-feira, das 9h às 17h, horário do Pacífico).

Criar um cluster do Google Kubernetes Engine

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

Criar um pool de nós da TPU

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

Em execução sem o modo privilegiado

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

Execute cargas de trabalho em pools de nós da 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 de TPU, especifique dois seletores para cada nó no manifesto do Google Kubernetes Engine:

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

As cargas de trabalho de treinamento e as cargas de trabalho de inferência seções 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 de 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.

Múltiplas fatias: ir além de uma única fatia

Você pode agregar frações menores em uma multislice para lidar com partes maiores cargas de trabalho de treinamento. Para mais informações, consulte Multifrações do Cloud TPU.

Tutoriais de carga de trabalho de treinamento

O foco destes tutoriais é o treinamento de cargas de trabalho em uma fração de TPU de vários hosts (para exemplo, 4 máquinas v5e). Elas abrangem os seguintes modelos:

Modelos Hugging Face FLAX: treine Diffusion no Pokémon
PyTorch/XLA: GPT2 no WikiText

Fazer o download dos recursos do tutorial

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

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

Criar e conectar ao cluster

Criar um cluster regional do GKE para que o Kubernetes é replicado em três zonas, o que proporciona maior disponibilidade. Crie seu cluster em us-west4, us-east1 ou us-central2, dependendo do qual A versão da 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 está 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 Identidade da carga de trabalho e Os recursos do driver CSI do Cloud Storage FUSE no cluster porque o exemplo as cargas de trabalho de inferência neste guia usam buckets do Cloud Storage para armazenar os 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 os recursos de driver CSI do Cloud Storage FUSE e Identidade da carga de trabalho 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 Hugging Face FLAX: treine Diffusion no Pokémon

Este exemplo treina o modelo de difusão estável do HuggingFace usando a Pokémon (em inglês) no conjunto de dados.

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

Criar imagem Docker

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

Antes de executar o comando a seguir, verifique se sua conta tem a configuração correta permissões para o Docker enviar 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 carga de trabalho

Crie um arquivo com o seguinte conteúdo 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-o 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 em PyTorch/XLA usando o conjunto de dados wikitext.

Criar imagem Docker

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

Antes de executar o comando a seguir, verifique se sua conta tem a configuração correta permissões para o Docker enviar 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 carga de trabalho

Copie o YAML a seguir e salve-o em um arquivo chamado tpu_job_gpt.yaml. Preencher 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

Neste tutorial, mostramos 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 o PyTorch. De modo geral, há quatro etapas separadas a serem realizadas no Cluster do GKE:

Criar um bucket do Cloud Storage e configurar o acesso a ele. Você usa um O bucket do Cloud Storage é usado para armazenar o modelo pré-treinado.
Fazer o download e converter um modelo pré-treinado em um modelo compatível com TPU. Aplique um Pod do Kubernetes que faz o download do modelo pré-treinado, usa O Cloud TPU Converter e armazena os modelos convertidos em um Cloud Storage usando o driver CSI do Cloud Storage FUSE. O Cloud TPU Converter e não exige hardware especializado. Neste tutorial, mostramos como fazer o download o modelo e execute o Cloud TPU Converter no pool de nós da CPU.
Inicie o servidor do modelo convertido. Aplique uma implantação. que exibe o modelo usando um framework de servidor apoiado pelo volume armazenado em o volume permanente ReadOnlymuitos (ROX, na sigla em inglês). As réplicas de implantação precisam ser executadas em um pool de nós de fração v5e com um pod do Kubernetes por nó. em um pool de nós de fração v5e com um pod do Kubernetes por nó.
Implantar um balanceador de carga para testar o servidor de modelo. O servidor está exposto solicitações externas usando o serviço LoadBalancer; Um script Python é fornecido com um exemplo de solicitação para testar a servidor de modelo.

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

Um diagrama que mostra o roteamento do balanceador de carga

Exemplos de implantação de servidor

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

O cluster está em execução com um pool de nós da 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 de TPU é 1

O manifesto do GKE a seguir define um único host do Google Cloud.

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ê usa um número diferente de nós no pool de nós da TPU, altere o replicas ao 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 Execute 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 ícones para o tipo de máquina correspondente.

Configuração

Faça o download dos scripts Python e dos manifestos YAML do tutorial usando os seguintes 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 Python usados neste tutorial exigem a versão 3.9 ou superior do Python. Lembre-se de instalar o requirements.txt para cada tutorial antes de executar a Scripts de teste do Python.

Se você não tiver a configuração apropriada do Python em seu ambiente local, será possível use o Cloud Shell para fazer o download e executar o scripts Python neste tutorial.

Configurar o cluster

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

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

As cargas de trabalho de exemplo pressupõem o seguinte:

Seu cluster está em execução com um pool de nós da 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, você precisa editar o manifesto de implantação do servidor.

Para a demonstração do JAX Stable Diffusion, você precisará de um pool de nós de CPU com Tipo de máquina com mais de 16 Gi+ de memória disponível (por exemplo, e2-standard-4). Isso é configurado no comando gcloud container clusters create ou pelo adicionar mais um pool de nós ao cluster atual 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 do nós a serem criados no pool.

Configurar o armazenamento de modelos

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

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

Execute o comando a seguir 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 Cloud Storage do Google Cloud.
your-bucket-location: o local do seu do bucket do Cloud Storage. Para melhorar o desempenho, especifique o local seu cluster do GKE está em execução.

Use as etapas a seguir para conceder ao cluster do GKE acesso aos do bucket. Para simplificar a configuração, os exemplos a seguir usam o e a conta de serviço padrão do Kubernetes. Para mais detalhes, consulte Configure o acesso a buckets do Cloud Storage usando a Identidade da carga de trabalho do GKE.

Crie uma conta de serviço do IAM para seu aplicativo ou use uma conta de serviço do IAM atual. Você pode 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 nesta etapa.
```
gcloud iam service-accounts create your-iam-service-acct \
--project=your-bucket-project-id
```
Substitua:
- your-iam-service-acct: o nome do novo serviço de IAM. do Compute Engine.
- your-bucket-project-id: o ID do projeto em que você criou o conta de serviço do IAM. A conta de serviço do IAM precisa estar no mesmo projeto que o bucket do Cloud Storage.
Verifique se a conta de serviço do IAM tem os papéis de armazenamento que você precisam.
```
gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
--member "serviceAccount:your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com" \
--role "roles/storage.objectAdmin"
```
Substitua:
- your-bucket-name: o nome do bucket do Cloud Storage.
- your-iam-service-acct: o nome do novo serviço de IAM. do Compute Engine.
- your-bucket-project-id: o ID do projeto em que você criou o conta de serviço do IAM.
Permitir que a conta de serviço do Kubernetes represente o IAM de conta de serviço adicionando 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 que atuam como a conta de serviço do IAM.
```
gcloud iam service-accounts add-iam-policy-binding your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:your-project-id.svc.id.goog[default/default]"
```
Substitua:
- your-iam-service-acct: o nome do novo serviço de IAM. do Compute Engine.
- your-bucket-project-id: o ID do projeto em que você criou o conta de serviço do IAM.
- your-project-id: o ID do projeto em que você criou o do cluster do GKE. Os buckets do Cloud Storage O cluster do GKE pode estar no mesmo projeto ou em projetos diferentes.
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=your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com
```
Substitua:
- your-iam-service-acct: o nome do novo serviço de IAM. do Compute Engine.
- your-bucket-project-id: o ID do projeto em que você criou o conta de serviço do IAM.
Execute o comando a seguir para preencher o nome do bucket nos arquivos YAML do nesta 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.
Crie a declaração de volume permanente e de volume permanente com o seguinte comando:
```
kubectl apply -f pvc-pv.yaml
```

Inferência e disponibilização de modelos JAX

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

pip install -r jax/requirements.txt

Execute a demonstração do serviço JAX BERT E2E:

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

O pod do Kubernetes executa as seguintes etapas:

Faz o download e usa o script Python export_bert_model.py do exemplo para fazer o download do modelo bert pré-treinado para um diretório temporário.
Usa a imagem do Cloud TPU Converter para converter o modelo pré-treinado de CPU para TPU e armazena o modelo no bucket do Cloud Storage que você criado durante a configuração.

Este pod do Kubernetes está configurado para ser executado na CPU do pool de nós padrão. Execute o 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 o STATUS ler o Completed.

Iniciar o servidor de modelo do TF para o modelo

As cargas de trabalho de exemplo neste tutorial pressupõem o seguinte:

O cluster está em execução com um pool de nós da 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, você precisa editar o manifesto de implantação do servidor.

Aplicar implantação

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

Verifique se o servidor está sendo executado com o seguinte comando:

kubectl get deployment bert-deployment

Pode levar um minuto para o AVAILABLE ler 3.

Aplicar serviço de balanceador de carga

kubectl apply -f jax/bert/loadbalancer.yaml

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

kubectl get svc tf-bert-service

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

Envie a solicitação ao servidor do modelo

Receba o IP externo do serviço de 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 exibição do JAX Stable Diffusion E2E

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

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

Exportar os modelos estáveis de difusão exige que o cluster tenha um nó de CPU com um tipo de máquina com mais de 16 GB de memória disponível, conforme descrito em Configure o cluster.

O pod do Kubernetes executa as seguintes etapas:

Faz o download e usa o script Python export_stable_diffusion_model.py no os recursos de exemplo para baixar o modelo de difusão estável pré-treinado para em um diretório temporário.
Usa a imagem do Cloud TPU Converter para converter o modelo pré-treinado de CPU para TPU e armazena o modelo no bucket do Cloud Storage criado por você durante a configuração do armazenamento.

Este pod do Kubernetes está configurado para ser executado no pool de nós de CPU padrão. Execute o 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 o STATUS ler o Completed.

Iniciar o contêiner do servidor de modelo 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 de TPU é 1

Se você estiver usando uma configuração de cluster diferente da descrita anteriormente, você precisa editar 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á sendo executado conforme o esperado:

kubectl get deployment stable-diffusion-deployment

Pode levar um minuto para o AVAILABLE ler 3.

Aplique o serviço de balanceador de carga:

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

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

kubectl get svc tf-stable-diffusion-service

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

Envie 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 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á salvo como stable_diffusion_images.jpg no seu diretório atual.

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 de tutorial em Python que enviam para o serviço de modelo do TF.

pip install -r tf/resnet50/requirements.txt

Etapa 1: converter o modelo

Aplicar 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 o STATUS ler o Completed.

Etapa 2: disponibilizar o modelo com o TensorFlow Serving

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

kubectl apply -f tf/resnet50/deployment.yml

Para verificar se o servidor está funcionando conforme esperado, use o seguinte comando:

kubectl get deployment resnet-deployment

Pode levar um minuto para o AVAILABLE ler 3.

Aplique o serviço de balanceador de carga:

kubectl apply -f tf/resnet50/loadbalancer.yml

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

kubectl get svc resnet-service

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

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

Consulte 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 do modelo.

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

A resposta será semelhante a esta:

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 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 do GKE e cluster quando terminar de usá-las.

Inferência e disponibilização de modelos PyTorch

Instale as dependências do Python para executar scripts de tutorial em Python que enviam para o serviço do modelo PyTorch:

pip install -r pt/densenet161/requirements.txt

Execute a demonstração de exibição do TorchServe Densenet161 E2E:

Gere 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 o STATUS ler o Completed.
Exiba o modelo com o TorchServe:
1. Aplique a implantação da disponibilização de modelos:
```
kubectl apply -f pt/densenet161/deployment.yml
```
2. Para verificar se o servidor está funcionando conforme esperado, use o seguinte comando:
```
kubectl get deployment densenet161-deployment
```
  Pode levar um minuto para o AVAILABLE ler 3.
3. Aplique o serviço de balanceador de carga:
```
kubectl apply -f pt/densenet161/loadbalancer.yml
```
  Verifique se o balanceador de carga está pronto para o tráfego externo com o seguinte comando:
```
kubectl get svc densenet161-service
```
  Pode levar alguns minutos para que EXTERNAL_IP tenha um IP listado.

Envie a solicitação de teste para o servidor do modelo:

Receba o IP externo do balanceador de carga:

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

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

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

Você verá uma resposta como esta:

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

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 do GKE e cluster quando você não precisa mais delas.

Como solucionar problemas comuns

Você pode encontrar informações sobre solução de problemas do GKE em Solução de problemas da TPU no GKE.

Falha na inicialização da TPU

Se você encontrar o erro a seguir, verifique se está executando a TPU. contêiner no modo privilegiado ou tiver aumentado o ulimit no contêiner do Docker. Para saber mais, consulte 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, A e B, e ambos sejam programados na TPU frações com uma determinada topologia de TPU (por exemplo, v4-32). Suponha também que você tem duas frações de TPU v4-32 no cluster do GKE; vamos chamaremos essa fração de X e de Y. Como seu cluster tem ampla capacidade para agendar os dois trabalhos, em teoria ambos devem ser agendados rapidamente – um trabalho em cada uma das duas frações v4-32 de TPU.

No entanto, sem um planejamento cuidadoso, é possível entrar em um cronograma um impasse. Suponha que o programador do Kubernetes programe um pod do Kubernetes a partir do job A na fração X e, em seguida, programa um pod do Kubernetes do job B na fração X. Em Nesse caso, considerando as regras de afinidade de pod do Kubernetes para o job A, o programador tentativa de programar todos os pods do Kubernetes restantes para o job A na fração X. O mesmo para Emprego B. Assim, nem o job A nem o job B poderão ser totalmente programados em uma uma única fração. O resultado será um impasse de programação.

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

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 do pool de nós da TPU com o Terraform

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

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

Se você tiver um cluster em que quer criar uma TPU de vários hosts pool de nós, é possível usar o seguinte snippet do Terraform:

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: o projeto do Google Cloud em que a carga de trabalho está sendo executada.
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 a carga de trabalho está sendo executada.
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 atual

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: o projeto do Google Cloud em que a carga de trabalho está sendo executada.
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 a carga de trabalho está sendo executada.