Gestisci un LLM utilizzando TPU Trillium su GKE con vLLM

Questo tutorial mostra come eseguire il servizio di modelli linguistici di grandi dimensioni (LLM) utilizzando le unità di elaborazione tensoriale (TPU) su Google Kubernetes Engine (GKE) con il framework di servizio vLLM. In questo tutorial, pubblica Llama 3.1 70b, utilizza TPU Trillium e configura la scalabilità automatica dei pod orizzontali utilizzando le metriche del server vLLM.

Questo documento è un buon punto di partenza se hai bisogno del controllo granulare, della scalabilità, della resilienza, della portabilità e dell'economicità di Kubernetes gestito quando esegui il deployment e il servizio dei tuoi workload di AI/ML.

Sfondo

Utilizzando TPU Trillium su GKE, puoi implementare una soluzione di pubblicazione affidabile e pronta per la produzione con tutti i vantaggi di Kubernetes gestito, inclusa una scalabilità efficiente e una maggiore disponibilità. Questa sezione descrive le tecnologie chiave utilizzate in questa guida.

TPU Trillium

Le TPU sono circuiti integrati per applicazioni specifiche (ASIC) sviluppati da Google. Le TPU vengono utilizzate per accelerare i modelli di machine learning e AI creati utilizzando framework come TensorFlow, PyTorch e JAX. Questo tutorial utilizza TPU Trillium, la TPU di sesta generazione di Google.

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

  1. Scopri di più sull'architettura di sistema di TPU Trillium.
  2. Scopri di più sulle TPU in GKE.

vLLM

vLLM è un framework open source altamente ottimizzato per l'erogazione di LLM. vLLM può aumentare il throughput di erogazione sulle TPU, con funzionalità come le seguenti:

  • Implementazione ottimizzata dei transformer con PagedAttention.
  • Raggruppamento continuo per migliorare il throughput complessivo della pubblicazione.
  • Parallelismo di tensori e pubblicazione distribuita su più TPU.

Per scoprire di più, consulta la documentazione di vLLM.

Cloud Storage FUSE

Cloud Storage FUSE fornisce l'accesso dal tuo cluster GKE a Cloud Storage per i pesi del modello che si trovano nei bucket di archiviazione di oggetti. In questo tutorial, il bucket Cloud Storage creato sarà inizialmente vuoto. All'avvio di vLLM, GKE scarica il modello da Hugging Face e memorizza nella cache i pesi nel bucket Cloud Storage. Al riavvio del pod o all'aumento di scala del deployment, i caricamenti successivi dei modelli scaricheranno i dati memorizzati nella cache dal bucket Cloud Storage, sfruttando i download paralleli per prestazioni ottimali.

Per saperne di più, consulta la documentazione del driver CSI di Cloud Storage FUSE.

Obiettivi

Questo tutorial è rivolto a DevOps o MLOps engineer o amministratori della piattaforma che vogliono utilizzare le funzionalità di orchestrazione di GKE per pubblicare LLM.

Questo tutorial illustra i seguenti passaggi:

  1. Crea un cluster GKE con la topologia TPU Trillium consigliata in base alle caratteristiche del modello.
  2. Esegui il deployment del framework vLLM su un pool di nodi nel cluster.
  3. Utilizza il framework vLLM per pubblicare Llama 3.1 70b utilizzando un bilanciatore del carico.
  4. Configura scalabilità automatica orizzontale dei pod utilizzando le metriche del server vLLM.
  5. Pubblica il modello.

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/iam.securityAdmin, roles/artifactregistry.writer, roles/container.clusterAdmin

    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 Nuovi principali, 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 questa sezione esegui il provisioning delle risorse necessarie per eseguire il deployment di vLLM e del modello.

Ottieni l'accesso al modello

Devi firmare il contratto di consenso per l'utilizzo di Llama 3.1 70b nel repository Hugging Face.

Genera un token di accesso

Se non ne hai già uno, genera un nuovo token Hugging Face:

  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.

Avvia Cloud Shell

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

Per configurare l'ambiente con Cloud Shell:

  1. Nella Google Cloud console, 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 PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)") && \
export CLUSTER_NAME=CLUSTER_NAME && \
export ZONE=ZONE && \
export REGION=REGION && \
export HF_TOKEN=HUGGING_FACE_TOKEN && \
export CLUSTER_VERSION=CLUSTER_VERSION && \
export GSBUCKET=GSBUCKET && \
export KSA_NAME=KSA_NAME && \
export NAMESPACE=NAMESPACE

    Sostituisci i seguenti valori:

    • PROJECT_ID : il tuo Google Cloud ID progetto.
    • CLUSTER_NAME : il nome del tuo cluster GKE.
    • ZONE : una zona che supporta TPU Trillium (v6e).
    • REGION : una regione che supporta TPU Trillium (v6e).
    • CLUSTER_VERSION : la versione 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. TPU Trillium (v6e) è supportata nelle versioni GKE 1.31.2-gke.1115000 o successive.
    • GSBUCKET : il nome del bucket Cloud Storage da utilizzare per Cloud Storage FUSE.
    • KSA_NAME : il nome dell'account ServiceAccount Kubernetes utilizzato per accedere ai bucket Cloud Storage. L'accesso ai bucket è necessario per il funzionamento di Cloud Storage FUSE.
    • NAMESPACE : lo spazio dei nomi Kubernetes in cui vuoi eseguire il deployment delle risorse vLLM.

Crea un cluster GKE

Puoi pubblicare LLM 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

  1. Crea un cluster GKE Autopilot:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
    --cluster-version=${CLUSTER_VERSION} \
    --region=${REGION}

Standard

  1. Crea un cluster GKE Standard:

    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --cluster-version=${CLUSTER_VERSION} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons GcsFuseCsiDriver

  2. Crea un pool di nodi di slice TPU:

    gcloud container node-pools create tpunodepool \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --num-nodes=1 \
    --machine-type=ct6e-standard-8t \
    --cluster=${CLUSTER_NAME} \
    --enable-autoscaling --total-min-nodes=1 --total-max-nodes=2

    GKE crea le seguenti risorse per l'LLM:

Configura kubectl per comunicare con il cluster

Per configurare kubectl in modo che comunichi con il cluster, esegui il seguente comando:

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

Crea un secret di Kubernetes per le credenziali di Hugging Face

  1. Crea uno spazio dei nomi. Puoi saltare questo passaggio se utilizzi lo spazio dei nomi default:

    kubectl create namespace ${NAMESPACE}

  2. Crea un secret di Kubernetes contenente il token di Hugging Face, esegui il seguente comando:

    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --namespace ${NAMESPACE}

Crea un bucket Cloud Storage

In Cloud Shell, esegui questo comando:

gcloud storage buckets create gs://${GSBUCKET} \
    --uniform-bucket-level-access

Viene creato un bucket Cloud Storage per archiviare i file del modello che hai scaricato da Hugging Face.

Configura un account di servizio Kubernetes per accedere al bucket

  1. Crea l'account di servizio Kubernetes:

    kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}

  2. Concedi l'accesso in lettura/scrittura all'account di servizio Kubernetes per accedere al bucket Cloud Storage:

    gcloud storage buckets add-iam-policy-binding gs://${GSBUCKET} \
  --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
  --role "roles/storage.objectUser"

  3. In alternativa, puoi concedere l'accesso in lettura e scrittura a tutti i bucket Cloud Storage del progetto:

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
--member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
--role "roles/storage.objectUser"

    GKE crea le seguenti risorse per l'LLM:

    1. Un bucket Cloud Storage per archiviare il modello scaricato e la cache di compilazione. Un driver CSI di Cloud Storage FUSE legge i contenuti del bucket.
    2. Volumi con la memorizzazione nella cache dei file abilitata e la funzionalità di download parallelo di Cloud Storage FUSE.
    Best practice:

    Utilizza una cache di file basata su tmpfs o Hyperdisk / Persistent Disk a seconda delle dimensioni previste dei contenuti del modello, ad esempio i file di peso. In questo tutorial utilizzerai la cache dei file di Cloud Storage FUSE basata su RAM.

Esegui il deployment del server di modelli vLLM

Per eseguire il deployment del server del modello vLLM, questo tutorial utilizza un deployment Kubernetes. Un deployment è un oggetto dell'API Kubernetes che ti consente di eseguire più repliche di pod distribuite tra i nodi di un cluster.

  1. Controlla il manifest del deployment salvato come vllm-llama3-70b.yaml:

    
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-tpu
  template:
    metadata:
      labels:
        app: vllm-tpu
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/cpu-limit: "0"
        gke-gcsfuse/memory-limit: "0"
        gke-gcsfuse/ephemeral-storage-limit: "0"
    spec:
      serviceAccountName: KSA_NAME
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
      containers:
      - name: vllm-tpu
        image: docker.io/vllm/vllm-tpu:73aa7041bfee43581314e6f34e9a657137ecc092
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --tensor-parallel-size=8
        - --max-model-len=4096
        - --model=meta-llama/Llama-3.1-70B
        - --download-dir=/data
        - --max-num-batched-tokens=512
        - --max-num-seqs=128
        env: 
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        - name: VLLM_XLA_CACHE_PATH
          value: "/data"
        - name: VLLM_USE_V1
          value: "1"
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: 8
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 15
          periodSeconds: 10
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
        - name: dshm
          mountPath: /dev/shm
      volumes:
      - name: gke-gcsfuse-cache
        emptyDir:
          medium: Memory
      - name: dshm
        emptyDir:
          medium: Memory
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: GSBUCKET
            mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    app: vllm-tpu
  type: LoadBalancer	
  ports:
    - name: http
      protocol: TCP
      port: 8000  
      targetPort: 8000

  2. Applica il manifest eseguendo il seguente comando:

    kubectl apply -f vllm-llama3-70b.yaml -n ${NAMESPACE}

  3. Visualizza i log del server del modello in esecuzione:

    kubectl logs -f -l app=vllm-tpu -n ${NAMESPACE}

    L'output dovrebbe essere simile al seguente:

    INFO:     Started server process [1]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

Pubblica il modello

  1. Per ottenere l'indirizzo IP esterno del servizio VLLM, esegui il seguente comando:

    export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}' -n ${NAMESPACE})

  2. Interagisci con il modello utilizzando curl:

    curl http://$vllm_service:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "meta-llama/Llama-3.1-70B",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
}'

    L'output dovrebbe essere simile al seguente:

    {"id":"cmpl-6b4bb29482494ab88408d537da1e608f","object":"text_completion","created":1727822657,"model":"meta-llama/Llama-3-8B","choices":[{"index":0,"text":" top holiday destination featuring scenic beauty and","logprobs":null,"finish_reason":"length","stop_reason":null,"prompt_logprobs":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}

Configurare l'autoscalatore personalizzato

In questa sezione, configuri scalabilità automatica orizzontale dei pod utilizzando le metriche Prometheus personalizzate. Utilizzi le metriche di Google Cloud Managed Service per Prometheus dal server vLLM.

Per saperne di più, consulta Google Cloud Managed Service per Prometheus. Questa opzione dovrebbe essere attiva per impostazione predefinita nel cluster GKE.

  1. Configura l'adattatore Stackdriver delle metriche personalizzate nel tuo cluster:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml

  2. Aggiungi il ruolo Visualizzatore di monitoraggio all'account di servizio utilizzato dall'adattatore Stackdriver per le metriche personalizzate:

    gcloud projects add-iam-policy-binding projects/${PROJECT_ID} \
    --role roles/monitoring.viewer \
    --member=principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter

  3. Salva il seguente manifest come vllm_pod_monitor.yaml:

    
apiVersion: monitoring.googleapis.com/v1
kind: PodMonitoring
metadata:
 name: vllm-pod-monitoring
spec:
 selector:
   matchLabels:
     app: vllm-tpu
 endpoints:
 - path: /metrics
   port: 8000
   interval: 15s

  4. Applicalo al cluster:

    kubectl apply -f vllm_pod_monitor.yaml -n ${NAMESPACE}

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 i seguenti comandi:

ps -ef | grep load.sh | awk '{print $2}' | xargs -n1 kill -9

gcloud container clusters delete ${CLUSTER_NAME} \
  --region=${REGION}

Passaggi successivi