Questa pagina è stata tradotta dall'API Cloud Translation.

Gestisci un LLM utilizzando TPU Trillium su GKE con vLLM

Questo tutorial mostra come gestire modelli linguistici di grandi dimensioni (LLM) utilizzando le unità di elaborazione tensoriale (TPU) su Google Kubernetes Engine (GKE) con il framework di gestione vLLM. In questo tutorial utilizzerai Llama 3.1 70b, TPU Trillium e configurerai la scalabilità automatica pod orizzontale 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 gestisci i tuoi carichi di lavoro AI/ML.

Sfondo

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

TPU Trillium

Le TPU sono circuiti integrati specifici per le applicazioni (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 la 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:

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

vLLM

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

Implementazione ottimizzata del transformer con PagedAttention.
Batching continuo per migliorare la velocità effettiva complessiva della pubblicazione.
Parallelismo dei tensori e servizio distribuito su più TPU.

Per saperne di più, consulta la documentazione su vLLM.

Nota: questo tutorial si concentra sul deployment di vLLM in una configurazione a singolo host, ideale per i modelli che possono essere pubblicati da una singola sezione TPU, come Llama 3.1 70b su un tipo di macchina ct6e-standard-8t. È importante notare che le configurazioni multi-host non sono supportate quando utilizzi vLLM con le TPU su GKE. La mancanza del supporto multi-host limita l'utilizzo di vLLM per la gestione di modelli estremamente grandi (ad esempio con oltre 400 miliardi di parametri) che richiedono la memoria e la potenza di calcolo aggregate di più host. Per i sistemi o i modelli di produzione che richiedono una configurazione multi-host, la soluzione consigliata e ottimizzata per il rendimento è utilizzare JetStream, il motore di Google per l'inferenza TPU. Per iniziare a utilizzare un deployment multi-host, consulta Inferenza JetStream MaxText su TPU v6e.

Cloud Storage FUSE

Cloud Storage FUSE fornisce l'accesso dal 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 inizialmente sarà 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 delle dimensioni del deployment, i carichi successivi del modello 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.

Crea un cluster GKE

Puoi gestire 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 workload, consulta Scegliere una modalità operativa GKE.

Autopilot

Crea un cluster GKE Autopilot:

gcloud container clusters create-auto ${CLUSTER_NAME} \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${CONTROL_PLANE_LOCATION}

Standard

Crea un cluster GKE Standard:

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

Crea un pool di nodi slice TPU:
```
gcloud container node-pools create tpunodepool \
    --location=${CONTROL_PLANE_LOCATION} \
    --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:
- Un cluster GKE Standard che utilizza Workload Identity Federation for GKE e in cui è abilitato il driver CSI di Cloud Storage FUSE.
- Un pool di nodi TPU Trillium con un tipo di macchina ct6e-standard-8t. Questo pool di nodi ha un nodo, otto chip TPU e la scalabilità automatica abilitata.

Configura kubectl per comunicare con il cluster

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

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

Crea un secret Kubernetes per le credenziali di Hugging Face

Crea uno spazio dei nomi. Puoi saltare questo passaggio se utilizzi lo spazio dei nomi default:
```
kubectl create namespace ${NAMESPACE}
```

Crea un secret Kubernetes che contenga il token Hugging Face eseguendo 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

In questo modo viene creato un bucket Cloud Storage per archiviare i file del modello che scarichi da Hugging Face.

Configura un service account Kubernetes per accedere al bucket

Crea il ServiceAccount Kubernetes:

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

Concedi l'accesso in lettura/scrittura al service account 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"

In alternativa, puoi concedere l'accesso in lettura e scrittura a tutti i bucket Cloud Storage nel 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 il contenuto 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 dei file supportata da tmpfs o Hyperdisk / Persistent Disk a seconda delle dimensioni previste dei contenuti del modello, ad esempio i file di peso. In questo tutorial, utilizzi la cache dei file Cloud Storage FUSE supportata dalla RAM.

Esegui il deployment del server del modello vLLM

Per eseguire il deployment del server del modello vLLM, 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.

Esamina il seguente manifest di Deployment salvato come vllm-llama3-70b.yaml, che utilizza una singola replica:


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: vllm/vllm-tpu:latest
        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

Se aumenti lo scale up del deployment a più repliche, le scritture simultanee in VLLM_XLA_CACHE_PATH causeranno l'errore: RuntimeError: filesystem error: cannot create directories. Per evitare questo errore, hai due opzioni:

Rimuovi la posizione della cache XLA rimuovendo il seguente blocco dal file YAML di deployment. Ciò significa che tutte le repliche ricompileranno la cache.
```
- name: VLLM_XLA_CACHE_PATH
  value: "/data"
```
Scala il deployment a 1 e attendi che la prima replica diventi pronta e scriva nella cache XLA. Poi esegui lo scale up a repliche aggiuntive. Ciò consente alle repliche rimanenti di leggere la cache senza tentare di scriverla.

Applica il manifest eseguendo questo comando:

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

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

Per ottenere l'indirizzo IP esterno del servizio VLLM, esegui questo comando:

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

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

Configura lo scalatore automatico personalizzato

In questa sezione configurerai la scalabilità automatica orizzontale dei pod utilizzando 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 abilitata per impostazione predefinita sul cluster GKE.

Configura l'adattatore Stackdriver delle metriche personalizzate sul tuo cluster:

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

Aggiungi il ruolo Visualizzatore Monitoring al 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
```
Nota: assicurati che il account di servizio utilizzato dal tuo cluster GKE disponga del ruolo Scrittore metriche di monitoraggio. Questo tutorial utilizza il service account predefinito di Compute Engine.

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

Applica il file al cluster:

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

Crea carico sull'endpoint vLLM

Crea carico sul server vLLM per testare la scalabilità automatica di GKE con una metrica vLLM personalizzata.

Esegui uno script bash (load.sh) per inviare N richieste parallele all'endpoint vLLM:

#!/bin/bash
N=PARALLEL_PROCESSES
export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}' -n ${NAMESPACE})
for i in $(seq 1 $N); do
  while true; do
    curl http://$vllm_service:8000/v1/completions -H "Content-Type: application/json" -d '{"model": "meta-llama/Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 1000, "temperature": 0}'
  done &  # Run in the background
done
wait

Sostituisci PARALLEL_PROCESSES con il numero di processi paralleli che vuoi eseguire.

Esegui lo script bash:
```
chmod +x load.sh
nohup ./load.sh &
```

Verifica che Google Cloud Managed Service per Prometheus acquisisca le metriche

Dopo che Google Cloud Managed Service per Prometheus ha eseguito lo scraping delle metriche e stai aggiungendo carico all'endpoint vLLM, puoi visualizzare le metriche in Cloud Monitoring.

Nella console Google Cloud , vai alla pagina Esplora metriche.

Vai a Esplora metriche
Fai clic su < > PromQL.
Inserisci la seguente query per osservare le metriche sul traffico:
```
vllm:num_requests_waiting{cluster='CLUSTER_NAME'}
```

Un grafico a linee mostra la metrica vLLM (num_requests_waiting) misurata nel tempo. La metrica vLLM aumenta da 0 (precaricamento) a un valore (post-caricamento). Questo grafico conferma che le metriche vLLM vengono importate in Google Cloud Managed Service per Prometheus. Il seguente grafico di esempio mostra un valore di precaricamento iniziale pari a 0, che raggiunge un valore massimo post-caricamento di quasi 400 entro un minuto.

vllm:num_requests_waiting

Esegui il deployment della configurazione di Horizontal Pod Autoscaler

Quando decidi su quale metrica eseguire la scalabilità automatica, ti consigliamo le seguenti metriche per la TPU vLLM:

num_requests_waiting: questa metrica si riferisce al numero di richieste in attesa nella coda del server del modello. Questo numero inizia a crescere in modo significativo quando la cache KV è piena.
gpu_cache_usage_perc: questa metrica si riferisce all'utilizzo della cache KV, che è direttamente correlato al numero di richieste elaborate per un determinato ciclo di inferenza sul server del modello. Tieni presente che questa metrica funziona allo stesso modo su GPU e TPU, anche se è legata allo schema di denominazione delle GPU.

Ti consigliamo di utilizzare num_requests_waiting quando esegui l'ottimizzazione per il throughput e il costo e quando i tuoi target di latenza sono raggiungibili con il throughput massimo del server del modello.

Ti consigliamo di utilizzare gpu_cache_usage_perc quando hai carichi di lavoro sensibili alla latenza in cui lo scaling basato sulla coda non è abbastanza veloce per soddisfare i tuoi requisiti.

Per ulteriori spiegazioni, consulta Best practice per la scalabilità automatica dei workload di inferenza del modello linguistico di grandi dimensioni (LLM) con le TPU.

Quando selezioni un target averageValue per la configurazione di HPA, devi determinarlo in modo sperimentale. Consulta il post del blog Risparmia sulle GPU: scalabilità automatica più intelligente per i carichi di lavoro di inferenza GKE per ulteriori idee su come ottimizzare questa parte. profile-generator utilizzato in questo post del blog funziona anche per vLLM TPU.

Nelle istruzioni riportate di seguito, esegui il deployment della configurazione HPA utilizzando la metrica num_requests_waiting. A scopo dimostrativo, imposta la metrica su un valore basso in modo che la configurazione HPA aumenti le repliche vLLM a due. Per eseguire il deployment della configurazione di Horizontal Pod Autoscaler utilizzando num_requests_waiting, segui questi passaggi:

Salva il seguente manifest come vllm-hpa.yaml:


apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
 name: vllm-hpa
spec:
 scaleTargetRef:
   apiVersion: apps/v1
   kind: Deployment
   name: vllm-tpu
 minReplicas: 1
 maxReplicas: 2
 metrics:
   - type: Pods
     pods:
       metric:
         name: prometheus.googleapis.com|vllm:num_requests_waiting|gauge
       target:
         type: AverageValue
         averageValue: 10

Le metriche vLLM in Google Cloud Managed Service per Prometheus seguono il formato vllm:metric_name.

Best practice:

Utilizza num_requests_waiting per scalare la velocità effettiva. Utilizza gpu_cache_usage_perc per i casi d'uso della TPU sensibili alla latenza.

Esegui il deployment della configurazione di Horizontal Pod Autoscaler:
```
kubectl apply -f vllm-hpa.yaml -n ${NAMESPACE}
```
GKE pianifica il deployment di un altro pod, il che attiva il gestore della scalabilità automatica pool di nodi per aggiungere un secondo nodo prima di eseguire il deployment della seconda replica vLLM.

Osserva l'avanzamento della scalabilità automatica dei pod:

kubectl get hpa --watch -n ${NAMESPACE}

L'output è simile al seguente:

NAME       REFERENCE             TARGETS       MINPODS   MAXPODS   REPLICAS   AGE
vllm-hpa   Deployment/vllm-tpu   <unknown>/10   1         2         0          6s
vllm-hpa   Deployment/vllm-tpu   34972m/10      1         2         1          16s
vllm-hpa   Deployment/vllm-tpu   25112m/10      1         2         2          31s
vllm-hpa   Deployment/vllm-tpu   35301m/10      1         2         2          46s
vllm-hpa   Deployment/vllm-tpu   25098m/10      1         2         2          62s
vllm-hpa   Deployment/vllm-tpu   35348m/10      1         2         2          77s

Attendi 10 minuti e ripeti i passaggi della sezione Verifica che Google Cloud Managed Service per Prometheus importi le metriche. Google Cloud Managed Service per Prometheus ora acquisisce le metriche da entrambi gli endpoint vLLM.