Introdução às TPUs no GKE
Os clientes do Google Kubernetes Engine (GKE) agora podem criar pools de nós do Kubernetes contendo pods TPU v4 e v5e. Um pod de TPU é um grupo de dispositivos de TPU conectados por interconexões de alta velocidade. Isso é diferente de um pod do Kubernetes, que é a menor unidade de computação implantável que pode ser criada e gerenciada no Kubernetes. Para cargas de trabalho que não exigem um pod de TPU completo, é possível usar um subconjunto de um pod de TPU completo chamado fração de TPU. Assim como os pods de TPU completos, cada dispositivo TPU em uma fração tem uma VM de TPU própria. Chamamos uma VM de TPU e o dispositivo conectado dela como nó host ou de TPU. Para mais informações sobre pods de TPU, consulte Arquitetura do sistema.
Como o termo Pod usado no contexto do GKE geralmente significa um pod do Kubernetes, para evitar qualquer confusão, sempre nos referiremos a um conjunto de um ou mais dispositivos TPU como uma fatia.
Quando você trabalha com o GKE, primeiro precisa 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 aceita dois tipos de pools de nós de TPU:
Pool de nós de fração de TPU de vários hosts
Um pool de nós de frações 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 meio de uma interconexão de alta velocidade (ICI, na sigla em inglês). Um pool de nós de frações de TPU de vários hosts é imutável. Depois 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 mais um nó do Kubernetes (VM da TPU) a ele. Para adicionar uma fração de TPU extra a um cluster do GKE, crie 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, todos os nós dela nã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 de TPU (nós) na fração, forçando todos os pods do Kubernetes na carga de trabalho a serem removidos. Quando todas as VMs de TPU na fração estiverem em execução, os pods do Kubernetes poderão ser programados nas VMs de TPU na nova fração.
O diagrama a seguir mostra um exemplo de uma fatia de TPU de vários hosts v5litepod-16 (v5e). Essa fração tem quatro VMs de TPU. Cada VM da TPU tem quatro chips da TPU v5e conectados a interconexões de alta velocidade (ICI), e cada chip da TPU v5e tem um TensorCore.
O diagrama a seguir mostra um cluster do GKE contendo uma fração de TPU v5litepod-16
(v5e) (topologia: 4x4) e uma fração v5litepod-8
(v5e) de TPU (topologia: 2x4):
Para um exemplo de execução de uma carga de trabalho em uma fração de TPU de vários hosts, consulte Executar carga de trabalho em uma fração de TPU de vários hosts.
Pools de nós de fração de TPU de host único
Um pool de nós de fração de host único é um pool de nós que contém uma ou mais
VMs de TPU independentes. Cada uma dessas VMs tem um dispositivo TPU conectado a ela. As VMs em um pool de nós de fração de host podem se comunicar pela rede de data center (DCN, na sigla em inglês), mas 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
:
Para um exemplo de execução de uma carga de trabalho em uma fração de TPU de host único, consulte Executar cargas de trabalho em nós de TPU.
Tipos de máquinas de TPU para pools de nós do GKE
Antes de criar pools de nós, você precisa escolher a versão e o tamanho da TPU da fração de TPU que sua carga de trabalho exige. A TPU v4 é compatível com a versão 1.26.1-gke.1500
e mais recentes do GKE Standard, a v5e na versão Standard do GKE 1.27.2-gke.2100
e posteriores, e a v5p na versão 1.28.3-gke.1024000
e mais recentes do GKE Standard.
As TPUs v4, v5e e v5p têm suporte na versão 1.29.2-gke.1521000
e mais recentes do Autopilot
do GKE.
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 fração de TPU (uma topologia de TPU) com base no tamanho do seu 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 frações.
V5e
Veja 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, 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 em 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 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. Para mais informações, consulte Tipos de aceleradores de treinamento v5e e Tipos de acelerador de inferência v5e.
As máquinas TPU v5e estão disponíveis nos idiomas us-west4-a
, us-east5-b
e us-east1-c
.
Os clusters GKE Standard precisam executar o plano de controle
versão 1.27.2-gke.2100 ou posterior.
O Autopilot do GKE precisa executar o plano de controle
versão 1.29.2-gke.1521000 ou mais recente.
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 | Alto | Média | Baixa |
Para liberar espaço para VMs com mais chips, o programador do GKE pode antecipar e reprogramar VMs com menos chips. Portanto, é provável que as VMs de 8 chips sejam interrompidas por 1 e 4 chips.
v4 e v5p
A seguir, estão 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 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áquina de pod da TPU v4 estão disponíveis em us-central2-b
. Os
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 fração de TPU v5p, use o tipo de máquina ct5p-hightpu-4t
,
que tem um host e quatro chips. Os tipos de máquina do pod TPU v5p estão disponíveis em us-west4-a
e us-east5-a
.
Os clusters do GKE Standard precisam
executar o plano de controle a versão 1.28.3-gke.1024000
ou mais recente.
O Autopilot do GKE precisa executar 1.29.2-gke.1521000
ou mais recente.
Para mais informações sobre a v5p, consulte Introdução ao treinamento da 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 da TPU.
- Somente reservas ESPECÍFICAS: ao usar TPUs no GKE,
SPECIFIC
é o único valor compatível com a sinalização--reservation-affinity
do comandogcloud container node-pools create
. - Somente a variante de VMs spot das TPUs preemptivas: as VMs spot são semelhantes às VMs preemptivas e estão sujeitas às mesmas limitações de disponibilidade, mas não têm 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 clusters pode calcular a capacidade incorretamente para novos nós da TPU antes que eles estejam disponíveis. O autoescalador de cluster pode realizar mais escalonar verticalmente e, como resultado, criar mais nós do que o necessário. O escalonador automático de cluster reduzirá os nós adicionais, se não forem necessários, após a operação normal de reduzir escala vertical.
- O escalonador automático cancela o escalonamento vertical: o escalonador automático de cluster cancela o escalonamento vertical de pools de nós de TPU que permanecem em status de espera por mais de 10 horas. O escalonador automático de cluster tentará realizar essas operações de escalonar verticalmente novamente mais tarde. Esse comportamento pode reduzir a capacidade de recebimento da TPU para clientes que não usam reservas.
- O Taint pode impedir a redução da escala vertical: as cargas de trabalho que não são de TPU e têm tolerância para o taint da TPU podem impedir reduzir escala vertical do pool de nós se forem recriadas durante a drenagem do pool de nós da TPU.
Garanta cotas suficientes de TPU e GKE
Talvez seja necessário aumentar determinadas cotas relacionadas ao GKE nas regiões onde seus recursos são criados.
As cotas a seguir têm valores padrão que provavelmente precisarão ser aumentados:
- Cota SSD (GB) de disco permanente: o disco de inicialização de cada nó do Kubernetes requer 100 GB por padrão. Portanto, essa cota precisa ser definida pelo menos até o número máximo de nós do GKE que você espera criar * 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ê a criação.
Para solicitar um aumento na 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ê tiver alguma dificuldade para aprovar suas solicitações de aumento de cota dentro de alguns dias, 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 um novo sistema de reservas baseado no Compute Engine.
Há várias informações importantes que você precisa saber sobre essa migração:
- A capacidade de TPU migrada para o novo sistema de reserva baseado no Compute Engine não pode ser usada com a API Queued Resource do Cloud TPU. Se você pretende usar recursos na fila da TPU com sua reserva, será necessário migrar apenas uma parte da reserva de TPU para o novo sistema de reservas baseado no Compute Engine.
- Nenhuma carga de trabalho pode ser executada ativamente nas 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 de conta do Google Cloud para programá-la. A janela de tempo de 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 na 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.
Como executar sem o modo privilegiado
Se quiser reduzir o escopo da permissão no contêiner, consulte Modo de privilégio da TPU.
Executar cargas de trabalho em nós da TPU
Consulte Executar cargas de trabalho em nós da TPU na documentação do Google Kubernetes Engine.
seletores de nodes
Para que o Kubernetes programe sua carga de trabalho em nós de TPU, é preciso especificar dois seletores para cada nó de TPU no manifesto do Google Kubernetes Engine:
- Defina
cloud.go888ogle.com.fqhub.com/gke-accelerator-type
comotpu-v5-lite-podslice
outpu-v4-podslice
. - Defina
cloud.go888ogle.com.fqhub.com/gke-tpu-topology
como a topologia de TPU do nó da TPU.
As seções Cargas de trabalho de treinamento e Cargas de trabalho de inferência contêm exemplos de manifestos que ilustram o uso desses seletores de nós.
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 de TPU
Se um nó da TPU em um pool de nós da fatia da TPU de vários hosts não estiver íntegro, todo o pool de nós será recriado. Para mais informações, consulte Reparo automático de nós na documentação do GKE.
Multislice: ir além de uma única fração
É possível agregar frações menores juntas em um multislice para lidar com cargas de trabalho de treinamento maiores. Para mais informações, consulte Cloud TPU Multislice.
Tutoriais de carga de trabalho de treinamento
Esses tutoriais se concentram nas cargas de trabalho de treinamento em uma fração de TPU de vários hosts (por exemplo, máquinas 4 v5e). Eles abrangem os seguintes modelos:
- Modelos FLAX: Treine a difusão no Pokémon
- PyTorch/XLA: GPT2 no WikiText
Fazer o download dos recursos do tutorial
Faça o download dos scripts Python do tutorial e das especificações YAML para cada modelo pré-treinado com o seguinte comando:
git clone http://github.com/GoogleCloudPlatform/ai-on-gke.git
Criar e conectar ao cluster
Criar um cluster regional do GKE para que o plano de controle
do Kubernetes seja replicado em três zonas, fornecendo maior disponibilidade.
Crie o cluster em us-west4
, us-east1
ou us-central2
, dependendo da versão da TPU usada. 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 os recursos da Identidade da carga de trabalho e 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 os recursos do driver CSI do Cloud Storage FUSE e da 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 seguinte comando para se conectar ao cluster recém-criado:
gcloud container clusters get-credentials cluster-name \ --location=cluster-region
Modelos FLAX: Treine a difusão no Pokémon
Este exemplo treina o modelo de difusão estável do HuggingFace usando o conjunto de dados Pokémon.
O modelo de difusão estável é um modelo latente de texto para imagem que gera imagens realistas a partir de qualquer entrada de texto. Para mais informações sobre difusão estável, consulte:
Criar imagem Docker
O Dockerfile está localizado na pasta
ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/
. Execute os comandos a seguir para criar e enviar a imagem 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 conteúdo a seguir e nomeie-o como tpu_job_diffusion.yaml
.
Preencha o campo 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.go888ogle.com.fqhub.com/gke-tpu-accelerator: tpu-v5-lite-podslice
cloud.go888ogle.com.fqhub.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 for executado, você poderá excluí-lo usando:
kubectl delete -f tpu_job_diffusion.yaml
PyTorch/XLA: GPT2 no WikiText
Neste tutorial, mostramos como executar o GPT2 em TPUs v5e usando 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/
.
Execute os comandos a seguir para criar e enviar a imagem 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 em um arquivo chamado tpu_job_gpt.yaml
. Preencha o campo 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.go888ogle.com.fqhub.com/gke-tpu-accelerator: tpu-v5-lite-podslice
cloud.go888ogle.com.fqhub.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 PyTorch. De modo geral, há quatro etapas separadas a serem executadas no cluster do GKE:
Crie um bucket do Cloud Storage e configure o acesso a ele. Use um bucket do Cloud Storage 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 faça o download do modelo pré-treinado, use o Cloud TPU Converter e armazene os modelos convertidos em um bucket do Cloud Storage usando o driver CSI do Cloud Storage FUSE. O conversor do Cloud TPU 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 da CPU.
Inicie o servidor do modelo convertido. Aplique uma implantação que disponibilize o modelo usando um framework de servidor apoiado pelo volume armazenado no volume permanente ReadOnly compatibilidades (ROX, na sigla em inglês). As réplicas de implantação precisam ser executadas em um nó de TPU do pod v5e com um pod do Kubernetes por nó.
Implantar 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 uma solicitação de exemplo para testar o servidor do modelo.
O diagrama a seguir mostra como as solicitações são roteadas pelo 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 é 1 x 1
- o número de chips de TPU é 1
O manifesto do GKE a seguir define uma única implantação do 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.go888ogle.com.fqhub.com/gke-tpu-topology: 1x1 # target topology
cloud.go888ogle.com.fqhub.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 da TPU, altere 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 seu manifesto:
spec:
securityContext:
privileged: true
Você não precisa executar os pods do Kubernetes no modo privilegiado na versão 1.28 ou posterior do GKE. Para mais detalhes, consulte Executar contêineres sem o modo privilegiado.
Se você estiver usando outro tipo de máquina:
- Defina
cloud.go888ogle.com.fqhub.com/gke-tpu-topology
como a topologia do tipo de máquina que você está usando. - Defina os dois campos
google.com/tpu
emresources
para que correspondam ao número de ícones do tipo de máquina correspondente.
Configuração
Faça o download dos scripts Python e dos manifestos YAML do tutorial usando o seguinte comando:
git clone http://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 os scripts de teste do Python.
Se você não tiver a configuração correta do Python no ambiente local, poderá usar o Cloud Shell para fazer o download e executar os scripts do 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
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, será necessário editar o manifesto de implantação do servidor.
Para a demonstração do JAX Stable Diffusion, você precisará de um pool de nós da CPU com um
tipo de máquina que tenha 16 Gi+ 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 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 dos nós a serem criados no pool de nós.
Configurar o armazenamento do modelo
Há várias maneiras de armazenar o modelo para disponibilização. Neste tutorial, usaremos a seguinte abordagem:
- Para converter o modelo pré-treinado para funcionar em TPUs, usaremos uma nuvem privada virtual apoiada por um Persistent Disk com acesso
ReadWriteMany
(RWX). - Para disponibilizar o modelo em várias TPUs de host único, usaremos 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 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á sendo executado.
Use as etapas a seguir para conceder ao cluster do GKE acesso ao bucket. 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 detalhes, consulte Configurar 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. É possível usar qualquer conta de serviço do IAM no projeto do bucket do Cloud Storage.
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.
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: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 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.
Permita que a conta de serviço do Kubernetes falsifique a conta de serviço do IAM 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 atue 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 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.
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 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.
Execute o comando a seguir 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.Crie o volume permanente e a declaração de volume permanente com o seguinte comando:
kubectl apply -f pvc-pv.yaml
Inferência e exibição de modelo JAX
Instale dependências do Python para executar scripts Python de tutorial que enviam solicitações ao serviço de modelo JAX.
pip install -r jax/requirements.txt
Execute a demonstração de exibição do JAX BERT E2E:
Essa demonstração usa um modelo BERT pré-treinado da Hugging Face.
O pod do Kubernetes executa as seguintes etapas:
- 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 para um diretório temporário. - Usa a imagem do conversor do Cloud TPU para converter o modelo pré-treinado de CPU para TPU e o armazena no bucket do Cloud Storage criado durante a setup.
Esse pod do Kubernetes está configurado para ser executado na CPU do pool de nós padrão. 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 o STATUS
ler Completed
.
Iniciar o servidor do modelo 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, será necessário 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:
kubectl get deployment bert-deployment
Pode levar um minuto para 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 da seguinte maneira:
kubectl get svc tf-bert-service
Pode levar alguns minutos para que 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 exibição do JAX Stable Diffusion E2E
Esta demonstração usa o modelo de difusão estável pré-treinado da Hugging Face.
Exportar modelo salvo do TF2 compatível com TPU do modelo de difusão estável do Flax
Exportar os 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 Gi+ de memória disponível, conforme descrito em Configurar cluster.
O pod do Kubernetes executa as seguintes etapas:
- Faz o download e usa o script Python
export_stable_diffusion_model.py
dos recursos de exemplo para fazer o download do modelo de difusão estável pré-treinado para um diretório temporário. - Usa a imagem do conversor do Cloud TPU para converter o modelo pré-treinado de CPU para TPU e o armazena no bucket do Cloud Storage criado durante a configuração de armazenamento.
Esse pod do Kubernetes está configurado para ser executado no pool de nós da 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 o STATUS
ler Completed
.
Iniciar o contêiner do servidor do 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 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 é 1 x 1
- o número de chips de TPU é 1
Se você estiver usando uma configuração de cluster diferente da descrita anteriormente, será necessário 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 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 da seguinte maneira:
kubectl get svc tf-stable-diffusion-service
Pode levar alguns minutos para que 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 script para enviar uma solicitação ao servidor
python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP
Saída esperada:
O prompt é Painting of a squirrel skating in New York
e a imagem de saída
será salva 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 dependências do Python para executar scripts Python de tutorial que enviam solicitações ao serviço de modelo do TF.
pip install -r tf/resnet50/requirements.txt
Etapa 1: converter o modelo
Aplique a conversão do 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 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
Para verificar se o servidor está sendo executado conforme o esperado, use o seguinte comando:
kubectl get deployment resnet-deployment
Pode levar um minuto para 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 da seguinte maneira:
kubectl get svc resnet-service
Pode levar alguns minutos para que EXTERNAL_IP
tenha um IP listado.
Etapa 3: enviar uma solicitação de teste ao 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 uma 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 e o cluster do GKE quando terminar de usá-los.
Inferência e exibição de modelos do PyTorch
Instale as dependências do Python para executar scripts Python do tutorial que enviam solicitações ao serviço de modelo PyTorch:
pip install -r pt/densenet161/requirements.txt
Execute a demonstração de veiculação do TorchServe Densenet161 E2E:
Gerar arquivo do modelo.
- Aplique o arquivo do modelo:
kubectl apply -f pt/densenet161/model-archive.yml
- Verifique se o modelo foi instalado corretamente com o seguinte:
kubectl get pods densenet161-model-archive
Pode levar alguns minutos para o
STATUS
lerCompleted
.Exiba o modelo com o TorchServe:
Aplique a implantação de disponibilização do modelo:
kubectl apply -f pt/densenet161/deployment.yml
Para verificar se o servidor está sendo executado conforme o esperado, use o seguinte comando:
kubectl get deployment densenet161-deployment
Pode levar um minuto para
AVAILABLE
ler3
.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.
Enviar solicitação de teste para o servidor do modelo:
Consiga 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 uma 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 e o cluster do GKE quando terminar de usá-los.
Como solucionar problemas comuns
Confira informações sobre solução de problemas do GKE em Resolver problemas de TPU no GKE.
Falha na inicialização da TPU
Se você encontrar o erro a seguir, verifique se está executando o contêiner de TPU no modo privilegiado ou se aumentou o ulimit
dentro do contêiner. Para mais informações, consulte Execução sem 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
Impasse de programação
Suponha que você tenha dois jobs, Job A e Job B, que serão programados em frações de TPU com 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 X e Y. Como seu cluster tem ampla capacidade para programar os dois jobs, em teoria, ambos precisam ser programados rapidamente: um em cada uma das duas frações v4-32
da TPU.
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 fração X e, em seguida, programe um pod do Kubernetes do job B na fração X. Nesse caso, considerando as regras de afinidade de pod do Kubernetes para o job A, o programador tentará programar todos os pods do Kubernetes restantes para o job A na fração X. O mesmo vale para o job B. Assim, nem o Job A nem o Job B poderão ser totalmente programados em uma única fração. O resultado será um impasse de programação.
Para evitar o risco de um impasse de programação, é possível usar a antiafinidade do pod
do Kubernetes com cloud.go888ogle.com.fqhub.com/gke-nodepool
como 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.go888ogle.com.fqhub.com/gke-nodepool
podAntiAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: job
operator: NotIn
values:
- pi
topologyKey: cloud.go888ogle.com.fqhub.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 de 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 atual do GKE
Se você tiver um cluster em que quer criar um pool de nós de TPU de vários hosts, use 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 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 atual 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
: o 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.