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

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.

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.

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"}