Gestisci modelli open source utilizzando le TPU su GKE con Optimum TPU

Questo tutorial mostra come pubblicare modelli open source di modelli linguistici di grandi dimensioni (LLM) utilizzando le Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con il framework di pubblicazione Optimum TPU di Hugging Face. In questo tutorial, scaricherai modelli open source da Hugging Face ed eseguirai il deployment dei modelli su un cluster GKE Standard utilizzando un container che esegue Optimum TPU.

Questa guida fornisce un punto di partenza se hai bisogno del controllo granulare, della scalabilità, della resilienza, della portabilità e dell'economicità di Kubernetes gestito durante il deployment e la gestione dei tuoi carichi di lavoro AI/ML.

Questo tutorial è destinato ai clienti dell'AI generativa nell'ecosistema Hugging Face, agli utenti nuovi o esistenti di GKE, agli ingegneri ML, agli ingegneri MLOps (DevOps) o agli amministratori della piattaforma interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per il servizio di LLM.

Ti ricordiamo che hai diverse opzioni per l'inferenza LLM su Google Cloud, che includono offerte come Vertex AI, GKE e Google Compute Engine, dove puoi incorporare librerie di serving come JetStream, vLLM e altre offerte dei partner. Ad esempio, puoi utilizzare JetStream per ottenere le ultime ottimizzazioni dal progetto. Se preferisci le opzioni di Hugging Face, puoi utilizzare Optimum TPU.

Optimum TPU supporta le seguenti funzionalità:

  • Batching continuo
  • Streaming dei token
  • Ricerca greedy e campionamento multinomiale utilizzando i transformer.

Ottenere l'accesso al modello

Puoi utilizzare i modelli Gemma 2B o Llama3 8B. Questo tutorial si concentra su questi due modelli, ma Optimum TPU supporta altri modelli.

Gemma 2B

Per accedere ai modelli Gemma per il deployment su GKE, devi prima firmare il contratto di consenso alla licenza, quindi generare un token di accesso Hugging Face.

Per utilizzare Gemma devi firmare il contratto di consenso. Segui queste istruzioni:

  1. Accedi alla pagina del consenso del modello.
  2. Verifica il consenso utilizzando il tuo account Hugging Face.
  3. Accetta i termini del modello.

Generare un token di accesso

Genera un nuovo token Hugging Face se non ne hai già uno:

  1. Fai clic su Il tuo profilo > Impostazioni > Token di accesso.
  2. Fai clic su Nuovo token.
  3. Specifica un nome a tua scelta e un ruolo di almeno Read.
  4. Fai clic su Generate a token (Genera un token).
  5. Copia il token generato negli appunti.

Llama3 8B

Per utilizzare Llama3 8b nel repository Hugging Face, devi firmare il contratto di consenso.

Generare un token di accesso

Genera un nuovo token Hugging Face se non ne hai già uno:

  1. Fai clic su Il tuo profilo > Impostazioni > Token di accesso.
  2. Seleziona Nuovo token.
  3. Specifica un nome a tua scelta e un ruolo di almeno Read.
  4. Seleziona Genera un token.
  5. Copia il token generato negli appunti.

Crea un cluster GKE

Crea un cluster GKE Standard con un nodo CPU:

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

Crea un pool di nodi TPU

Crea un pool di nodi TPU v5e con 1 nodo e 8 chip:

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

Se le risorse TPU sono disponibili, GKE esegue il provisioning del pool di nodi. Se le risorse TPU non sono temporaneamente disponibili, l'output mostra un messaggio di errore GCE_STOCKOUT. Per risolvere i problemi relativi all'esaurimento delle TPU, consulta Risorse TPU insufficienti per soddisfare la richiesta di TPU.

Configura kubectl per comunicare con il cluster:

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

Crea il container

Esegui il comando make per creare l'immagine

cd optimum-tpu && make tpu-tgi

Esegui il push dell'immagine in 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

Crea un secret Kubernetes per le credenziali di Hugging Face

Crea un secret di Kubernetes che contenga il token Hugging Face:

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

Implementa Optimum TPU

Per eseguire il deployment di Optimum TPU, questo tutorial utilizza un deployment Kubernetes. Un deployment è un oggetto API Kubernetes che ti consente di eseguire più repliche di pod distribuite tra i nodi di un cluster.

Gemma 2B

  1. Salva il seguente manifest di deployment come 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

    Questo manifest descrive un deployment di TPU Optimum con un bilanciatore del carico interno sulla porta TCP 8080.

  2. Applica il manifest

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

Llama3 8B

  1. Salva il seguente manifest come 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

    Questo manifest descrive un deployment di TPU Optimum con un bilanciatore del carico interno sulla porta TCP 8080.

  2. Applica il manifest

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

Visualizza i log del deployment in esecuzione:

kubectl logs -f -l app=tgi-tpu

L'output dovrebbe essere simile al seguente:

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

Assicurati che il modello sia stato scaricato completamente prima di procedere alla sezione successiva.

Pubblica il modello

Configura il port forwarding sul modello:

kubectl port-forward svc/service 8080:8080

Interagisci con il server del modello utilizzando curl

Verifica i modelli di cui è stato eseguito il deployment:

In una nuova sessione del terminale, utilizza curl per chattare con il modello:

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'

L'output dovrebbe essere simile al seguente:

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