Gestisci un LLM utilizzando TPU su GKE con JetStream e PyTorch

Questa guida mostra come gestire un modello linguistico di grandi dimensioni (LLM) utilizzando le Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con JetStream tramite PyTorch. In questa guida, scaricherai i pesi del modello in Cloud Storage ed eseguirai il deployment su un cluster GKE Autopilot o Standard utilizzando un container che esegue JetStream.

Se hai bisogno della scalabilità, della resilienza e dell'efficacia in termini di costi offerti dalle funzionalità di Kubernetes durante il deployment del modello su JetStream, questa guida è un buon punto di partenza.

Questa guida è destinata ai clienti dell'AI generativa che utilizzano PyTorch, 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 la gestione degli LLM.

Sfondo

Se gestisci un LLM utilizzando le TPU su GKE con JetStream, puoi creare una soluzione di gestione solida e pronta per la produzione con tutti i vantaggi di Kubernetes gestito, tra cui efficienza dei costi, scalabilità e maggiore disponibilità. Questa sezione descrive le tecnologie chiave utilizzate in questo tutorial.

Informazioni sulle TPU

Le TPU sono circuiti integrati specifici per le applicazioni (ASIC) sviluppati da Google e utilizzati per accelerare i modelli di machine learning e AI creati utilizzando framework come TensorFlow, PyTorch e JAX.

Prima di utilizzare le TPU in GKE, ti consigliamo di completare il seguente percorso di apprendimento:

  1. Scopri di più sulla disponibilità della versione attuale della TPU con l'architettura di sistema di Cloud TPU.
  2. Scopri di più sulle TPU in GKE.

Questo tutorial illustra l'erogazione di vari modelli LLM. GKE esegue il deployment del modello su nodi TPUv5e a singolo host con topologie TPU configurate in base ai requisiti del modello per la gestione dei prompt con bassa latenza.

Informazioni su JetStream

JetStream è un framework open source per la pubblicazione di inferenze sviluppato da Google. JetStream consente un'inferenza ad alte prestazioni, ad alta velocità effettiva e ottimizzato per la memoria su TPU e GPU. JetStream offre ottimizzazioni avanzate delle prestazioni, tra cui il batching continuo, l'ottimizzazione della cache KV e tecniche di quantizzazione, per facilitare l'implementazione degli LLM. JetStream consente l'erogazione di PyTorch/XLA e JAX TPU per ottenere prestazioni ottimali.

Raggruppamento continuo

Il batch continuo è una tecnica che raggruppa dinamicamente le richieste di inferenza in entrata in batch, riducendo la latenza e aumentando la velocità effettiva.

Quantizzazione della cache KV

La quantizzazione della cache KV prevede la compressione della cache chiave-valore utilizzata nei meccanismi di attenzione, riducendo i requisiti di memoria.

Quantizzazione dei pesi Int8

La quantizzazione dei pesi Int8 riduce la precisione dei pesi del modello da virgola mobile a 32 bit a numeri interi a 8 bit, il che comporta un calcolo più rapido e un utilizzo della memoria ridotto.

Per scoprire di più su queste ottimizzazioni, consulta i repository dei progetti JetStream PyTorch e JetStream MaxText.

Informazioni su PyTorch

PyTorch è un framework di machine learning open source sviluppato da Meta e ora parte della Linux Foundation. PyTorch fornisce funzionalità di alto livello come il calcolo dei tensori e le reti neurali profonde.

Obiettivi

  1. Prepara un cluster GKE Autopilot o Standard con la topologia TPU consigliata in base alle caratteristiche del modello.
  2. Esegui il deployment dei componenti JetStream su GKE.
  3. Recupera e pubblica il modello.
  4. Pubblica e interagisci con il modello pubblicato.

Architettura

Questa sezione descrive l'architettura GKE utilizzata in questo tutorial. L'architettura include un cluster GKE Autopilot o Standard che esegue il provisioning delle TPU e ospita i componenti JetStream per eseguire il deployment e pubblicare i modelli.

Il seguente diagramma mostra i componenti di questa architettura:

Architettura del cluster GKE con pool di nodi TPU a singolo host contenenti i componenti JetStream-PyTorch e JetStream HTTP.

Questa architettura include i seguenti componenti:

  • Un cluster regionale GKE Autopilot o Standard.
  • Due node pool di sezioni TPU a host singolo che ospitano il deployment di JetStream.
  • Il componente Service distribuisce il traffico in entrata a tutte le repliche JetStream HTTP.
  • JetStream HTTP è un server HTTP che accetta le richieste come wrapper per il formato richiesto di JetStream e le invia al client GRPC di JetStream.
  • JetStream-PyTorch è un server JetStream che esegue l'inferenza con il batch continuo.

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/resourcemanager.projectIamAdmin

    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 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.
  • Verifica di disporre di una quota sufficiente per otto chip TPU v5e PodSlice Lite. In questo tutorial utilizzi le istanze on demand.
  • Crea un token Hugging Face se non ne hai già uno.

Ottenere l'accesso al modello

Accedere a vari modelli su Hugging Face per il deployment su GKE

Gemma 7B-it

Per accedere al modello Gemma per il deployment su GKE, devi prima firmare il contratto di consenso alla licenza.

  1. Accedere alla pagina del consenso del modello Gemma su Hugging Face
  2. Accedi a Hugging Face, se non l'hai ancora fatto.
  3. Leggi e accetta i Termini e condizioni modello.

Llama 3 8B

Per accedere al modello Llama 3 per il deployment su GKE, devi prima firmare il contratto di consenso alla licenza.

  1. Accedere alla pagina di consenso del modello Llama 3 su Hugging Face
  2. Accedi a Hugging Face, se non l'hai ancora fatto.
  3. Leggi e accetta i Termini e condizioni modello.

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 consoleGoogle 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
gcloud config set billing/quota_project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export LOCATION=LOCATION
export CLUSTER_VERSION=CLUSTER_VERSION
export BUCKET_NAME=BUCKET_NAME

    Sostituisci i seguenti valori:

    • PROJECT_ID: il tuo Google Cloud ID progetto.
    • CLUSTER_NAME: il nome del tuo cluster GKE. Non è necessario specificare il prefisso gs://.
    • REGION: la regione in cui si trovano il cluster GKE 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). Per i cluster Autopilot, assicurati di disporre di risorse di zona TPU v5e sufficienti per la regione che preferisci.
    • (Solo cluster Standard) LOCATION: 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.
    • CLUSTER_VERSION: la versione di GKE, che deve supportare il tipo di macchina che vuoi utilizzare. Tieni presente che la versione GKE predefinita potrebbe non essere disponibile per la TPU di destinazione. Per un elenco delle versioni GKE minime disponibili per tipo di macchina TPU, consulta Disponibilità delle TPU in GKE.
    • BUCKET_NAME: il nome del bucket Cloud Storage utilizzato per archiviare la cache di compilazione JAX.

Creare e configurare risorse Google Cloud

Segui queste istruzioni per creare le risorse richieste.

Crea un cluster GKE

Puoi pubblicare Gemma su TPU in un cluster GKE Autopilot o Standard. Ti consigliamo di utilizzare un cluster Autopilot per un'esperienza Kubernetes completamente gestita. Per scegliere la modalità operativa GKE più adatta ai tuoi carichi di lavoro, consulta Scegliere una modalità operativa GKE.

Autopilot

Crea un cluster GKE Autopilot:

gcloud container clusters create-auto CLUSTER_NAME \
    --project=PROJECT_ID \
    --region=REGION \
    --cluster-version=CLUSTER_VERSION

Sostituisci CLUSTER_VERSION con la versione del cluster appropriata. Per un cluster GKE Autopilot, utilizza una versione del canale di rilascio regolare.

Standard

  1. Crea un cluster GKE Standard regionale che utilizza Workload Identity Federation for GKE:

    gcloud container clusters create CLUSTER_NAME \
    --enable-ip-alias \
    --machine-type=e2-standard-4 \
    --num-nodes=2 \
    --cluster-version=CLUSTER_VERSION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=REGION

    La creazione del cluster potrebbe richiedere diversi minuti.

    Sostituisci CLUSTER_VERSION con la versione del cluster appropriata.

  2. Crea un node pool TPU v5e con una topologia 2x4 e due nodi:

    gcloud container node-pools create tpu-nodepool \
  --cluster=CLUSTER_NAME \
  --machine-type=ct5lp-hightpu-8t \
  --project=PROJECT_ID \
  --num-nodes=2 \
  --region=REGION \
  --node-locations=LOCATION

Genera il token della CLI di Hugging Face in Cloud Shell

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. Modifica le autorizzazioni del token di accesso per avere l'accesso in lettura al repository Hugging Face del modello.
  6. Copia il token generato negli appunti.

Crea un secret di Kubernetes per le credenziali di Hugging Face

In Cloud Shell:

  1. Configura kubectl per comunicare con il cluster:

    gcloud container clusters get-credentials CLUSTER_NAME --location=REGION

  2. Crea un secret per archiviare le credenziali di Hugging Face:

    kubectl create secret generic huggingface-secret \
    --from-literal=HUGGINGFACE_TOKEN=HUGGINGFACE_TOKEN

    Sostituisci HUGGINGFACE_TOKEN con il tuo token Hugging Face.

Configura l'accesso ai tuoi workload utilizzando Workload Identity Federation for GKE

Assegna un ServiceAccount Kubernetes all'applicazione e configura questo service account Kubernetes in modo che funga da account di servizio IAM.

  1. Crea un account di servizio IAM per la tua applicazione:

    gcloud iam service-accounts create wi-jetstream

  2. Aggiungi un'associazione di policy IAM per il tuo account di servizio IAM per gestire Cloud Storage:

    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member "serviceAccount:wi-jetstream@PROJECT_ID.iam.gserviceaccount.com" \
    --role roles/storage.objectUser

gcloud projects add-iam-policy-binding PROJECT_ID \
    --member "serviceAccount:wi-jetstream@PROJECT_ID.iam.gserviceaccount.com" \
    --role roles/storage.insightsCollectorService

  3. Consenti al service account Kubernetes di rappresentare l'account di servizio IAM aggiungendo un'associazione dei criteri IAM tra i due service account. Questa associazione consente al service account Kubernetes di agire come account di servizio IAM:

    gcloud iam service-accounts add-iam-policy-binding wi-jetstream@PROJECT_ID.iam.gserviceaccount.com \
    --role roles/iam.workloadIdentityUser \
    --member "serviceAccount:PROJECT_ID.svc.id.goog[default/default]"

  4. Annota il account di servizio Kubernetes con l'indirizzo email del account di servizio IAM:

    kubectl annotate serviceaccount default \
    iam.gke.io/gcp-service-account=wi-jetstream@PROJECT_ID.iam.gserviceaccount.com

Esegui il deployment di JetStream

Esegui il deployment del container JetStream per pubblicare il modello. Questo tutorial utilizza i manifest di deployment di Kubernetes. Un deployment è un oggetto API Kubernetes che consente di eseguire più repliche di pod distribuite tra i nodi di un cluster.

Salva il seguente manifest come jetstream-pytorch-deployment.yaml:

Gemma 7B-it

apiVersion: apps/v1
kind: Deployment
metadata:
  name: jetstream-pytorch-server
spec:
  replicas: 2
  selector:
    matchLabels:
      app: jetstream-pytorch-server
  template:
    metadata:
      labels:
        app: jetstream-pytorch-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: jetstream-pytorch-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pytorch-server:v0.2.4
        args:
        - --model_id=google/gemma-7b-it
        - --override_batch_size=30
        - --enable_model_warmup=True
        volumeMounts:
        - name: huggingface-credentials
          mountPath: /huggingface
          readOnly: true
        ports:
        - containerPort: 9000
        resources:
          requests:
            google.com/tpu: 8
          limits:
            google.com/tpu: 8
        startupProbe:
          httpGet:
            path: /healthcheck
            port: 8000
            scheme: HTTP
          periodSeconds: 60
          initialDelaySeconds: 90
          failureThreshold: 50
        livenessProbe:
          httpGet:
            path: /healthcheck
            port: 8000
            scheme: HTTP
          periodSeconds: 60
          failureThreshold: 30
        readinessProbe:
          httpGet:
            path: /healthcheck
            port: 8000
            scheme: HTTP
          periodSeconds: 60
          failureThreshold: 30
      - name: jetstream-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
        ports:
        - containerPort: 8000
      volumes:
      - name: huggingface-credentials
        secret:
          defaultMode: 0400
          secretName: huggingface-secret
---
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pytorch-server
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000

Llama 3 8B

apiVersion: apps/v1
kind: Deployment
metadata:
  name: jetstream-pytorch-server
spec:
  replicas: 2
  selector:
    matchLabels:
      app: jetstream-pytorch-server
  template:
    metadata:
      labels:
        app: jetstream-pytorch-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: jetstream-pytorch-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pytorch-server:v0.2.4
        args:
        - --model_id=meta-llama/Meta-Llama-3-8B
        - --override_batch_size=30
        - --enable_model_warmup=True
        volumeMounts:
        - name: huggingface-credentials
          mountPath: /huggingface
          readOnly: true
        ports:
        - containerPort: 9000
        resources:
          requests:
            google.com/tpu: 8
          limits:
            google.com/tpu: 8
        startupProbe:
          httpGet:
            path: /healthcheck
            port: 8000
            scheme: HTTP
          periodSeconds: 60
          initialDelaySeconds: 90
          failureThreshold: 50
        livenessProbe:
          httpGet:
            path: /healthcheck
            port: 8000
            scheme: HTTP
          periodSeconds: 60
          failureThreshold: 30
        readinessProbe:
          httpGet:
            path: /healthcheck
            port: 8000
            scheme: HTTP
          periodSeconds: 60
          failureThreshold: 30
      - name: jetstream-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
        ports:
        - containerPort: 8000
      volumes:
      - name: huggingface-credentials
        secret:
          defaultMode: 0400
          secretName: huggingface-secret
---
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pytorch-server
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000

Il manifest imposta le seguenti proprietà chiave:

  • model_id: il nome del modello di Hugging Face (google/gemma-7b-it, meta-llama/Meta-Llama-3-8B) (vedi i modelli supportati).
  • override_batch_size: la dimensione del batch di decodifica per dispositivo, dove un chip TPU equivale a un dispositivo. Per impostazione predefinita, questo valore è 30.
  • enable_model_warmup: questa impostazione attiva il preriscaldamento del modello dopo l'avvio del server del modello. Per impostazione predefinita, questo valore è False.

Facoltativamente, puoi impostare queste proprietà:

  • max_input_length: la lunghezza massima della sequenza di input. Per impostazione predefinita, questo valore è 1024.
  • max_output_length: la lunghezza massima di decodifica dell'output. Il valore predefinito è 1024.
  • quantize_weights: indica se il checkpoint è quantizzato. Il valore predefinito è 0; impostalo su 1 per attivare la quantizzazione int8.
  • internal_jax_compilation_cache: la directory per la cache di compilazione JAX. Il valore predefinito è ~/jax_cache; impostalo su gs://BUCKET_NAME/jax_cache per la memorizzazione nella cache remota.

Nel manifest, un probe di avvio è configurato per garantire che il server del modello sia etichettato Ready dopo che il modello è stato caricato e il warmup è stato completato. I probe di attività e idoneità sono configurati per garantire l'integrità del server del modello.

  1. Applica il manifest:

    kubectl apply -f jetstream-pytorch-deployment.yaml

  2. Verifica il deployment:

    kubectl get deployment

    L'output è simile al seguente:

    NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
jetstream-pytorch-server          0/2     2            0           ##s

    Per i cluster Autopilot, il provisioning delle risorse TPU richieste potrebbe richiedere alcuni minuti.

  3. Visualizza i log del server JetStream-PyTorch per verificare che i pesi del modello siano stati caricati e che il warmup del modello sia stato completato. Il completamento di questa operazione potrebbe richiedere alcuni minuti.

    kubectl logs deploy/jetstream-pytorch-server -f -c jetstream-pytorch-server

    L'output è simile al seguente:

    Started jetstream_server....
2024-04-12 04:33:37,128 - root - INFO - ---------Generate params 0 loaded.---------

  4. Verifica che il deployment sia pronto:

    kubectl get deployment

    L'output è simile al seguente:

    NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
jetstream-pytorch-server          2/2     2            2           ##s

    La registrazione dell'endpoint healthcheck potrebbe richiedere alcuni minuti.

Pubblica il modello

In questa sezione, interagisci con il modello.

Configurare il port forwarding

Puoi accedere al deployment di JetStream tramite il servizio ClusterIP che hai creato nel passaggio precedente. I servizi ClusterIP sono raggiungibili solo dall'interno del cluster. Pertanto, per accedere al servizio dall'esterno del cluster, completa i seguenti passaggi:

Per stabilire una sessione di port forwarding, esegui questo comando:

kubectl port-forward svc/jetstream-svc 8000:8000

Interagisci con il modello utilizzando curl

  1. Verifica di poter accedere al server HTTP JetStream aprendo un nuovo terminale ed eseguendo questo comando:

    curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

    Il completamento della richiesta iniziale può richiedere diversi secondi a causa del riscaldamento del modello. L'output è simile al seguente:

    {
    "response": " for data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its readability, extensive libraries (pandas, scikit-learn), and integration with other tools.\n- High demand for Python programmers in data science roles.\n\n**2. R:**\n- Popular choice for data analysis and visualization, particularly in academia and research.\n- Extensive libraries for statistical modeling and data wrangling.\n\n**3. Java:**\n- Enterprise-grade platform for data science, with strong performance and scalability.\n- Widely used in data mining and big data analytics.\n\n**4. SQL:**\n- Essential for data querying and manipulation, especially in relational databases.\n- Used for data analysis and visualization in various industries.\n\n**5. Scala:**\n- Scalable and efficient for big data processing and machine learning models.\n- Popular in data science for its parallelism and integration with Spark and Spark MLlib."
}

Hai eseguito correttamente le seguenti operazioni:

  1. È stato eseguito il deployment del server di modelli JetStream-PyTorch su GKE utilizzando le TPU.
  2. Ha pubblicato e interagito con il modello.

Osserva le prestazioni del modello

Per osservare le prestazioni del modello, puoi utilizzare l'integrazione della dashboard JetStream in Cloud Monitoring. Con questa dashboard puoi visualizzare metriche sul rendimento critiche come il throughput dei token, la latenza delle richieste e i tassi di errore.

Per utilizzare la dashboard JetStream, devi abilitare Google Cloud Managed Service per Prometheus, che raccoglie le metriche da JetStream, nel tuo cluster GKE.

Puoi quindi visualizzare le metriche utilizzando la dashboard JetStream. Per informazioni sull'utilizzo di Google Cloud Managed Service per Prometheus per raccogliere le metriche dal modello, consulta le indicazioni sull'osservabilità di JetStream nella documentazione di Cloud Monitoring.

Risoluzione dei problemi

  • Se visualizzi il messaggio Empty reply from server, è possibile che il container non abbia completato il download dei dati del modello. Controlla di nuovo i log del pod per il messaggio Connected che indica che il modello è pronto per essere pubblicato.
  • Se visualizzi Connection refused, verifica che l'inoltro delle porte sia attivo.

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 create in questa guida, esegui i seguenti comandi e segui le istruzioni:

gcloud container clusters delete CLUSTER_NAME --region=REGION

gcloud iam service-accounts delete wi-jetstream@PROJECT_ID.iam.gserviceaccount.com

Passaggi successivi