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) 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 di 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 di piattaforme interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per la pubblicazione di LLM.

Ti ricordiamo che hai a disposizione 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 servizio 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 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

  • 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, roles/artifactregistry.admin

    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 column 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 l'accesso.

    4. Nel campo Nuove entità, inserisci il tuo identificatore dell'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.

      8. 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 , avvia una sessione di Cloud Shell facendo clic su Icona di attivazione di Cloud Shell Attiva Cloud Shell nella Google Cloud console. 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
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, 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 contribuire a garantire la disponibilità delle TPU, puoi utilizzare le prenotazioni TPU.

      Configura kubectl per comunicare con il tuo 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 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=ZONE

      Passaggi successivi