Pubblica modelli open source utilizzando TPU su GKE con Optimum TPU


Questo tutorial mostra come pubblicare modelli open source di modelli linguistici di grandi dimensioni (LLM), con Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con la TPU Optimum di pubblicazione di Hugging Face. In questo tutorial, puoi scaricare modelli open source da Hugging Face ed eseguire il deployment dei modelli su uno standard GKE 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 per il deployment e il servizio dei tuoi carichi di lavoro di IA/ML.

Questo tutorial è rivolto ai clienti dell'IA generativa nell'ecosistema Hugging Face, nuovi per gli utenti esistenti di GKE, ML Engineer, MLOps (DevOps) engineer amministratori di piattaforma interessati all'utilizzo di container di orchestrazione per la distribuzione degli LLM.

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

Optimum TPU supporta le seguenti funzionalità:

  • Raggruppamento continuo
  • Streaming di token
  • Ricerca greedy e campionamento multinomiale mediante 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

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

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

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

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

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

      Vai a IAM
    2. Seleziona il progetto.
    3. Fai clic su Concedi accesso.
    4. Nel campo Nuove entità, inserisci il tuo identificatore utente. In genere si tratta dell'indirizzo email di un Account Google.

    5. Nell'elenco Seleziona un ruolo, seleziona un ruolo.
    6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ogni ruolo aggiuntivo.
    7. Fai clic su Salva.

prepara l'ambiente

In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate in Google Cloud. Cloud Shell include il software di cui avrai bisogno per questo tutorial, tra cui kubectl e gcloud CLI.

Per configurare l'ambiente con Cloud Shell:

  1. Nella console Google Cloud, avvia una sessione di Cloud Shell facendo clic su Icona di attivazione di Cloud Shell Attiva Cloud Shell nella console Google Cloud. Viene avviata una sessione nel riquadro inferiore della console Google Cloud.

  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
    

    Sostituisci i seguenti valori:

    • PROJECT_ID: il tuo ID progetto Google Cloud.
    • CLUSTER_NAME: il nome del tuo cluster GKE.
    • REGION_NAME: la regione in cui GKE cluster, il bucket Cloud Storage e i nodi TPU. La regione contiene zone in cui sono disponibili 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.
  3. Clona il repository Optimum TPU:

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

Ottieni 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 in GKE, devi prima firmare il contratto di consenso per la licenza e poi generare un token di accesso a Hugging Face.

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

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

Genera un token di accesso

Genera un nuovo token Abbracciamento 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

Devi firmare il contratto di consenso per utilizzare Llama3 8b nel repository Hugging Face

Genera 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 1 nodo CPU:

gcloud container 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

Configura kubectl per comunicare con il cluster:

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

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 Kubernetes che contiene 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 -

Esegui il deployment di Optimum TPU

Esegui il deployment di TPU ottimale:

Gemma 2B

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

    Questo file 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=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
    

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

  2. Applicare 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 al modello:

kubectl port-forward svc/service 8080:8080

Interagire 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 creato in questa guida, esegui questo comando:

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

Passaggi successivi