Publique modelos de código aberto com TPUs no GKE com o Optimum TPU

Padrão

Este tutorial mostra como publicar modelos de código aberto de modelos de linguagem (conteúdo extenso) (MDI/CE) usando unidades de processamento tensor (TPUs) no Google Kubernetes Engine (GKE) com a estrutura de publicação Optimum TPU da Hugging Face. Neste tutorial, transfere modelos de código aberto do Hugging Face e implementa os modelos num cluster Standard do GKE através de um contentor que executa o Optimum TPU.

Este guia oferece um ponto de partida se precisar do controlo detalhado, da escalabilidade, da resiliência, da portabilidade e da rentabilidade do Kubernetes gerido ao implementar e publicar as suas cargas de trabalho de IA/ML.

Este tutorial destina-se a clientes de IA generativa no ecossistema do Hugging Face, utilizadores novos ou existentes do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas que tenham interesse em usar capacidades de orquestração de contentores do Kubernetes para publicar LLMs.

Lembramos que tem várias opções para a inferência de MDIs no Google Cloud, que abrangem ofertas como o Vertex AI, o GKE e o Google Compute Engine, onde pode incorporar bibliotecas de publicação como o JetStream, o vLLM e outras ofertas de parceiros. Por exemplo, pode usar o JetStream para receber as otimizações mais recentes do projeto. Se preferir as opções do Hugging Face, pode usar a TPU Optimum.

A TPU ideal suporta as seguintes funcionalidades:

• Processamento em lote contínuo
• Streaming de tokens
• Pesquisa gananciosa e amostragem multinomial com transformadores.

Objetivos

1. Prepare um cluster padrão do GKE com a topologia de TPU recomendada com base nas caraterísticas do modelo.
2. Implemente a TPU ideal no GKE.
3. Use a TPU ideal para publicar os modelos suportados através do curl.

Antes de começar

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

• Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin

  Check for the roles

  1. In the Google Cloud console, go to the IAM page.

     Go to IAM
  2. Select the project.

  3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

  4. For all rows that specify or include you, check the Role column to see whether the list of roles includes the required roles.

  Grant the roles

  1. In the Google Cloud console, go to the IAM page.

     Aceder ao IAM
  2. Selecione o projeto.
  3. Clique em Conceder acesso.

  4. No campo Novos responsáveis, introduza o identificador do utilizador. Normalmente, este é o endereço de email de uma Conta Google.

  5. Na lista Selecionar uma função, selecione uma função.
  6. Para conceder funções adicionais, clique em Adicionar outra função e adicione cada função adicional.
  7. Clique em Guardar.
  • Crie uma conta Hugging Face, se ainda não tiver uma.
  • Certifique-se de que o seu projeto tem quota suficiente para o Cloud TPU no GKE.

  Prepare o ambiente

  Neste tutorial, vai usar o Cloud Shell para gerir recursos alojados no Google Cloud. O Cloud Shell vem pré-instalado com o software de que precisa para este tutorial, incluindo o kubectl e a CLI gcloud.

  Para configurar o seu ambiente com o Cloud Shell, siga estes passos:

  1. Na Google Cloud consola, inicie uma sessão do Cloud Shell clicando em Ativar Cloud Shell na Google Cloud consola. Esta ação inicia uma sessão no painel inferior da consola. Google Cloud

  2. Defina as variáveis de ambiente predefinidas:

     gcloud config set project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export CLUSTER_NAME=CLUSTER_NAME
     export REGION=REGION_NAME
     export ZONE=ZONE
     export HF_TOKEN=HF_TOKEN
     

     Substitua os seguintes valores:

     • PROJECT_ID: o seu Google Cloud ID do projeto.
     • CLUSTER_NAME: o nome do seu cluster do GKE.
     • REGION_NAME: a região onde o cluster do GKE, o contentor do Cloud Storage e os nós da TPU estão localizados. A região contém zonas onde os tipos de máquinas TPU v5e estão disponíveis (por exemplo, us-west1, us-west4, us-central1, us-east1, us-east5 ou europe-west4).
     • (Apenas para clusters padrão) ZONE: a zona onde os recursos de TPU estão disponíveis (por exemplo, us-west4-a). Para clusters do Autopilot, não precisa de especificar a zona, apenas a região.
     • HF_TOKEN: o seu token da HuggingFace.

  3. Clone o repositório Optimum TPU:

     git clone https://github.com/huggingface/optimum-tpu.git
     

  Aceda ao modelo

  Pode usar os modelos Gemma 2B ou Llama3 8B. Este tutorial foca-se nestes dois modelos, mas o Optimum TPU suporta mais modelos.

  Gemma 2B

  Para aceder aos modelos Gemma para implementação no GKE, primeiro tem de assinar o contrato de consentimento de licença e, em seguida, gerar um token de acesso do Hugging Face.

  Assine o contrato de consentimento de licença

  Tem de assinar o contrato de consentimento para usar o Gemma. Siga estas instruções:

  1. Aceda à página de consentimento do modelo.
  2. Valide o consentimento através da sua conta do Hugging Face.
  3. Aceite os termos do modelo.

  Gere um token de acesso

  Gere um novo token do Hugging Face se ainda não tiver um:

  1. Clique em O seu perfil > Definições > Tokens de acesso.
  2. Clique em Novo token.
  3. Especifique um nome à sua escolha e uma função de, pelo menos, Read.
  4. Clique em Gerar um token.
  5. Copie o token gerado para a área de transferência.

  Llama3 8B

  Tem de assinar o contrato de consentimento para usar o Llama3 8b no Hugging Face Repo

  Gere um token de acesso

  Gere um novo token do Hugging Face se ainda não tiver um:

  1. Clique em O seu perfil > Definições > Tokens de acesso.
  2. Selecione Novo token.
  3. Especifique um nome à sua escolha e uma função de, pelo menos, Read.
  4. Selecione Gerar um token.
  5. Copie o token gerado para a área de transferência.

  Crie um cluster do GKE

  Crie um cluster padrão do GKE com 1 nó de CPU:

  gcloud container clusters create CLUSTER_NAME \
      --project=PROJECT_ID \
      --num-nodes=1 \
      --location=ZONE
  

  Crie um node pool de TPUs

  Crie um node pool de TPUs v5e com 1 nó e 8 chips:

  gcloud container node-pools create tpunodepool \
      --location=ZONE \
      --num-nodes=1 \
      --machine-type=ct5lp-hightpu-8t \
      --cluster=CLUSTER_NAME
  

  Se os recursos da TPU estiverem disponíveis, o GKE aprovisiona o conjunto de nós. Se os recursos da TPU estiverem temporariamente indisponíveis, o resultado mostra uma GCE_STOCKOUT mensagem de erro. Para resolver problemas de erros de indisponibilidade de TPUs, consulte o artigo Recursos de TPU insuficientes para satisfazer o pedido de TPU.

  Configure o kubectl para comunicar com o seu cluster:

  gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${ZONE}
  

  Crie o contentor

  Execute o comando make para criar a imagem

  cd optimum-tpu && make tpu-tgi
  

  Envie a imagem para o Artifact Registry

  gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
  gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
  docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
  docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
  

  Crie um segredo do Kubernetes para as credenciais do Hugging Face

  Crie um segredo do Kubernetes que contenha o token do Hugging Face:

  kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  

  Implemente a TPU Optimum

  Para implementar a TPU ideal, este tutorial usa uma implementação do Kubernetes. Uma implementação é um objeto da API Kubernetes que lhe permite executar várias réplicas de pods distribuídas entre os nós num cluster.

  Gemma 2B

  1. Guarde o seguinte manifesto de implementação como optimum-tpu-gemma-2b-2x4.yaml:

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: tgi-tpu
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: tgi-tpu
       template:
         metadata:
           labels:
             app: tgi-tpu
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 2x4
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           containers:
           - name: tgi-tpu
             image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
             args:
             - --model-id=google/gemma-2b
             - --max-concurrent-requests=4
             - --max-input-length=8191
             - --max-total-tokens=8192
             - --max-batch-prefill-tokens=32768
             - --max-batch-size=16
             securityContext:
                 privileged: true
             env:
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     name: hf-secret
                     key: hf_api_token
             ports:
             - containerPort: 80
             resources:
               limits:
                 google.com/tpu: 8
             livenessProbe:
               httpGet:
                 path: /health
                 port: 80
               initialDelaySeconds: 300
               periodSeconds: 120
     
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: service
     spec:
       selector:
         app: tgi-tpu
       ports:
         - name: http
           protocol: TCP
           port: 8080
           targetPort: 80
     

     Este manifesto descreve uma implementação de TPU ideal com um balanceador de carga interno na porta TCP 8080.

  2. Aplique o manifesto

     kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml
     

  Llama3 8B

  1. Guarde o seguinte manifesto como optimum-tpu-llama3-8b-2x4.yaml:

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: tgi-tpu
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: tgi-tpu
       template:
         metadata:
           labels:
             app: tgi-tpu
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 2x4
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           containers:
           - name: tgi-tpu
             image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
             args:
             - --model-id=meta-llama/Meta-Llama-3-8B
             - --max-concurrent-requests=4
             - --max-input-length=8191
             - --max-total-tokens=8192
             - --max-batch-prefill-tokens=32768
             - --max-batch-size=16
             env:
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     name: hf-secret
                     key: hf_api_token
             ports:
             - containerPort: 80
             resources:
               limits:
                 google.com/tpu: 8
             livenessProbe:
               httpGet:
                 path: /health
                 port: 80
               initialDelaySeconds: 300
               periodSeconds: 120
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: service
     spec:
       selector:
         app: tgi-tpu
       ports:
         - name: http
           protocol: TCP
           port: 8080
           targetPort: 80
     

     Este manifesto descreve uma implementação de TPU ideal com um balanceador de carga interno na porta TCP 8080.

  2. Aplique o manifesto

     kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml
     

  Veja os registos da implementação em execução:

  kubectl logs -f -l app=tgi-tpu
  

  O resultado deve ser semelhante ao seguinte:

  2024-07-09T22:39:34.365472Z  WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
  2024-07-09T22:40:47.851405Z  INFO text_generation_router: router/src/main.rs:314: Warming up model
  2024-07-09T22:40:54.559269Z  INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
  2024-07-09T22:40:54.559291Z  INFO text_generation_router: router/src/main.rs:352: Connected
  2024-07-09T22:40:54.559295Z  WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0
  

  Certifique-se de que o modelo é totalmente transferido antes de avançar para a secção seguinte.

  Publique o modelo

  Configure o encaminhamento de porta para o modelo:

  kubectl port-forward svc/service 8080:8080
  

  Interaja com o servidor de modelos através do curl

  Valide os modelos implementados:

  Numa nova sessão de terminal, use curl para conversar com o modelo:

  curl 127.0.0.1:8080/generate     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}'     -H 'Content-Type: application/json'
  

  O resultado deve ser semelhante ao seguinte:

  {"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}
  

  Limpar

  Para evitar incorrer em custos na sua conta do Google Cloud pelos recursos usados neste tutorial, elimine o projeto que contém os recursos ou mantenha o projeto e elimine os recursos individuais.

  Elimine os recursos implementados

  Para evitar incorrer em custos na sua Google Cloud conta pelos recursos que criou neste guia, execute o seguinte comando:

  gcloud container clusters delete CLUSTER_NAME \
    --location=ZONE
  

  O que se segue?

  • Explore a documentação da TPU ideal.
  • Descubra como pode executar modelos Gemma no GKE e como executar cargas de trabalho de IA/ML otimizadas com as capacidades de orquestração da plataforma GKE.
  • Saiba mais acerca das TPUs no GKE.