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

Padrão

---

Neste tutorial, mostramos como exibir modelos de código aberto e linguagem grande (LLM) que usam Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) com a TPU ideal de exibição do BigQuery do Hugging Face. Neste tutorial, você fará o download de modelos de código aberto do Hugging Face e os implantará em um cluster padrão do GKE usando um contêiner que executa a TPU ideal.

Este guia é um ponto de partida se você precisar do controle granular, da escalonabilidade, da resiliência, da portabilidade e da economia do Kubernetes gerenciado ao implantar e disponibilizar suas cargas de trabalho de IA/ML.

Este tutorial é destinado a clientes de IA generativa no ecossistema do Hugging Face, usuários novos ou existentes do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para veiculação de LLMs.

Como lembrete, você tem várias opções para inferência de LLM no Google Cloud, em ofertas como Vertex AI, GKE e Google Compute Engine, em que é possível incorporar bibliotecas de disponibilização como JetStream, vLLM e outras ofertas de parceiros. Por exemplo, você pode usar o JetStream para receber as otimizações mais recentes do projeto. Se você preferir as opções do Hugging Face, use a TPU ideal.

O TPU Optimum oferece suporte aos seguintes recursos:

• Lotes contínuos
• Streaming de token
• Greedy search e amostragem multinomial com transformadores.

Objetivos

1. Prepare um cluster do GKE Standard com a topologia de TPU recomendada com base nas características do modelo.
2. Implantar a TPU Optimum no GKE.
3. Use o Optimum TPU para disponibilizar os modelos compatíveis pelo curl.

Antes de começar

• Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

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

  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

• Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

• Enable the required API.

  Enable the API

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

  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

• Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

• Enable the required API.

  Enable the API

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

  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 colunn 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.

     Acessar o IAM
  2. Selecionar um projeto.
  3. Clique em CONCEDER ACESSO.

  4. No campo Novos principais, insira seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

  5. Na lista Selecionar um papel, escolha um.
  6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
  7. Clique em Salvar.

• Crie uma conta do Hugging Face caso ainda não tenha uma.
• Verifique se o projeto tem cota suficiente para a Cloud TPU no GKE.

Prepare o ambiente

Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e gcloud CLI.

Para configurar o ambiente com o Cloud Shell, siga estas etapas:

1. No console do Google Cloud, inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell no console do Google Cloud. Isso inicia uma sessão no painel inferior do console do Google Cloud.

2. Defina as variáveis de ambiente padrão:

   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
   

   Substitua os seguintes valores:

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

3. Clone o repositório do TPU ideal:

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

Receber acesso ao modelo

Você pode usar os modelos Gemma 2B ou Llama3 8B. Este tutorial se concentra nesses dois modelos, mas o Optimum TPU oferece suporte a mais modelos.

Gemma 2B

Para ter acesso aos modelos Gemma para implantação no GKE, primeiro assine o contrato de consentimento de licença e gere um token de acesso do Hugging Face.

Assinar o contrato de consentimento de licença

É necessário assinar o contrato de consentimento para usar o Gemma. Siga estas instruções:

1. Acesse a página de consentimento do modelo.
2. Confirme seu consentimento usando sua conta do Hugging Face.
3. Aceite os termos do modelo.

Gerar um token de acesso

Gere um novo token do Huggin' Face, caso ainda não tenha um:

1. Clique em Seu perfil > Configurações > Tokens de acesso.
2. Clique em Novo token.
3. Especifique um Nome de sua escolha e um Papel de pelo menos Read.
4. Clique em Generate a token.
5. Copie o token gerado para a área de transferência.

Llama3 8B

É necessário assinar o contrato de consentimento para usar o Llama3 8b no Hugging Face Repo

Gerar um token de acesso

Gere um novo token do Huggin' Face, caso ainda não tenha um:

1. Clique em Seu perfil > Configurações > Tokens de acesso.
2. Selecione Novo token.
3. Especifique um Nome de sua escolha e um Papel 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 do GKE Standard com um nó de CPU:

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

Criar um pool de nós de TPU

Crie um pool de nós de TPU 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

Configure o kubectl para se comunicar com o cluster:

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

Crie o contêiner

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

Criar um Secret do Kubernetes para as credenciais do Hugging Face

Crie um secret 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 -

Implantar a TPU ideal

Implantar a TPU ideal:

Gemma 2B

1. Salve o seguinte manifesto 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=32
           - --max-total-tokens=64
           - --max-batch-size=1
           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 implantação de TPU Optimum 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. Salve 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=32
           - --max-total-tokens=64
           - --max-batch-size=1
           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 implantação de TPU Optimum com um balanceador de carga interno na porta TCP 8080.

2. Aplique o manifesto.

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

Confira os registros da implantação em execução:

kubectl logs -f -l app=tgi-tpu

A saída será semelhante a esta:

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

Verifique se o download do modelo foi concluído antes de prosseguir para a próxima seção.

Disponibilizar o modelo

Configure o encaminhamento de portas para o modelo:

kubectl port-forward svc/service 8080:8080

Interagir com o servidor de modelo usando curl

Verifique os modelos implantados:

Em uma nova sessão do 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'

A saída será semelhante a esta:

{"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 cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir os recursos implantados

Para evitar cobranças na sua conta do Google Cloud pelos recursos criados neste guia, execute o seguinte comando:

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

A seguir

• Consulte a documentação do TPU Optimum.
• Descubra como executar modelos Gemma no GKE e como executar cargas de trabalho otimizadas de IA/ML com os recursos de orquestração da plataforma do GKE.
• Saiba mais sobre TPUs no GKE.