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

Standard

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, scarichi modelli open source da Hugging Face ed esegui 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.

I prodottiGoogle Cloud come GKE, Vertex AI e Compute Engine supportano varie librerie di pubblicazione, come JetStream, vLLM e altre offerte dei partner. Ad esempio, puoi utilizzare JetStream per ottenere le ultime ottimizzazioni del progetto. Se preferisci le opzioni di Hugging Face, puoi utilizzare Optimum TPU.

Optimum TPU supporta le seguenti funzionalità:

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

Obiettivi

1. Prepara un cluster GKE Standard con la topologia TPU consigliata in base alle caratteristiche del modello.
2. Esegui il deployment di Optimum TPU su GKE.
3. Utilizza Optimum TPU per gestire i modelli supportati tramite curl.

Prima di iniziare

Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.

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 role (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

If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

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 role (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

If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

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

• Crea un account Hugging Face, se non ne hai già uno.
• Assicurati che il tuo progetto disponga di una quota sufficiente per Cloud TPU in GKE.

Ruoli obbligatori

Per ottenere le autorizzazioni necessarie per configurare cluster e carichi di lavoro, chiedi all'amministratore di concederti i seguenti ruoli IAM sul progetto:

• Amministratore service account (roles/iam.serviceAccountAdmin)
• Gestisci cluster GKE: Amministratore Kubernetes Engine (roles/container.admin)
• Crea ed esegui il push delle immagini in Artifact Registry: Artifact Registry Administrator (roles/artifactregistry.admin)

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Prepara l'ambiente

In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate su Google Cloud. Cloud Shell include il software preinstallato necessario per questo tutorial, tra cui kubectl e gcloud CLI.

Per configurare l'ambiente con Cloud Shell:

1. Nella console Google Cloud , attiva Cloud Shell.

   Attiva Cloud Shell

   Nella parte inferiore della console Google Cloud viene avviata una sessione di Cloud Shell e viene visualizzato un prompt della riga di comando. Cloud Shell è un ambiente shell con Google Cloud CLI già installata e con valori già impostati per il progetto corrente. L'inizializzazione della sessione può richiedere alcuni secondi.

2. Imposta le variabili di ambiente predefinite:

   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
   

   Sostituisci i seguenti valori:

   • PROJECT_ID: il tuo Google Cloud ID progetto.
   • CLUSTER_NAME: il nome del tuo cluster GKE.
   • REGION_NAME: la regione in cui si trovano il cluster GKE, il bucket Cloud Storage e i nodi TPU. La regione contiene zone in cui sono disponibili i tipi di macchine TPU v5e (ad esempio us-west1, us-west4, us-central1, us-east1, us-east5 o europe-west4).
   • (Solo cluster Standard) ZONE: la zona in cui sono disponibili le risorse TPU (ad esempio, us-west4-a). Per i cluster Autopilot, non è necessario specificare la zona, ma solo la regione.
   • HF_TOKEN: il tuo token HuggingFace.

3. Clona il repository Optimum TPU:

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

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 e poi generare un token di accesso a Hugging Face.

Firmare il contratto di consenso alla licenza

Per utilizzare Gemma devi firmare l'accordo sul 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=REGION_NAME

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=REGION_NAME \
    --num-nodes=1 \
    --machine-type=ct5lp-hightpu-8t \
    --node-locations=ZONE \
    --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 alla disponibilità delle risorse, consulta Risorse TPU insufficienti per soddisfare la richiesta di TPU.

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

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina le risorse di cui è stato eseguito il deployment

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse che hai creato in questa guida, esegui questo comando:

gcloud container clusters delete CLUSTER_NAME \
  --location=REGION_NAME

Passaggi successivi

• Esplora la documentazione di Optimum TPU.
• Scopri come eseguire i modelli Gemma su GKE e come eseguire workload di AI/ML ottimizzati con le funzionalità di orchestrazione della piattaforma GKE.
• Scopri di più sulle TPU in GKE.