Questa guida mostra come gestire un modello linguistico di grandi dimensioni (LLM) Gemma utilizzando Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con JetStream tramite MaxText. In questa guida scaricherai in Cloud Storage i pesi del modello ottimizzato per l'istruzione con i parametri Gemma 7B e ne eseguirai il deployment su un cluster GKE Autopilot o Standard utilizzando un container che esegue JetStream.
Questa guida è un buon punto di partenza se hai bisogno delle funzionalità di scalabilità, resilienza ed economicità offerte dalle funzionalità di Kubernetes quando esegui il deployment di un modello su JetStream.
Contesto
Utilizzando Gemma mediante 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 in termini di costi, scalabilità e disponibilità superiore. Questa sezione descrive le tecnologie chiave utilizzate in questo tutorial.
Gemma
Gemma è un insieme di modelli di intelligenza artificiale generativa (AI), leggeri e a disponibilità aperta, rilasciati con licenza aperta. Questi modelli di IA possono essere eseguiti in applicazioni, hardware, dispositivi mobili o servizi in hosting. Puoi utilizzare i modelli Gemma per la generazione di testo, ma puoi anche ottimizzarli per attività specializzate.
Per saperne di più, consulta la documentazione di Gemma.
TPU
Le TPU sono ASIC (Application-Specific Integrated Circuit) appositamente sviluppati di Google per accelerare il machine learning e i modelli di 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:
- Scopri di più sulla disponibilità attuale delle versioni di TPU con l'architettura di sistema di Cloud TPU.
- Scopri di più sulle TPU in GKE.
Questo tutorial spiega la pubblicazione del modello Gemma 7B. GKE esegue il deployment del modello su nodi TPUv5e a host singolo con topologie TPU configurate in base ai requisiti del modello per la gestione di prompt con bassa latenza.
JetStream
JetStream è un framework di inferenza open source sviluppato da Google. JetStream consente un'inferenza ad alte prestazioni, elevata velocità effettiva e memoria ottimizzata su TPU e GPU. Offre ottimizzazioni avanzate delle prestazioni, tra cui tecniche di batch e quantizzazione continui, per facilitare il deployment dei modelli LLM. JetStream consente a PyTorch/XLA e alla gestione di JAX TPU di fornire prestazioni ottimali.
Per scoprire di più su queste ottimizzazioni, fai riferimento ai repository di progetto JetStream PyTorch e JetStream MaxText.
MaxText
MaxText è un'implementazione LLM JAX ad alte prestazioni, scalabile e adattabile, basata su librerie JAX open source come Flax, Orbax e Optax. L'implementazione LLM solo decoder di MaxText è scritta in Python. Sfrutta il compilatore XLA per ottenere prestazioni elevate senza dover creare kernel personalizzati.
Per saperne di più sugli ultimi modelli e dimensioni dei parametri supportati da MaxText, consulta il repository di MaxtText.
Obiettivi
Questo tutorial è destinato ai clienti di IA generativa che utilizzano JAX, a utenti nuovi o esistenti di GKE, ML engineer, ingegneri MLOps (DevOps) o amministratori di piattaforma interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per la gestione di LLM.
Questo tutorial illustra i seguenti passaggi:
- Prepara un cluster GKE Autopilot o Standard con la topologia TPU consigliata in base alle caratteristiche del modello.
- Eseguire il deployment dei componenti JetStream su GKE.
- Scarica e pubblica il modello ottimizzato per le istruzioni Gemma 7B.
- Pubblicare e interagire con il modello pubblicato.
Architettura
Questa sezione descrive l'architettura GKE utilizzata in questo tutorial. L'architettura comprende un cluster GKE Autopilot o Standard che esegue il provisioning delle TPU e ospita componenti JetStream per il deployment e la gestione dei modelli.
Il seguente diagramma mostra i componenti di questa architettura:
Questa architettura include i seguenti componenti:
- Un cluster GKE Autopilot o Standard a livello di regione.
- Due pool di nodi della sezione TPU con singolo host che ospitano il deployment JetStream.
- Il componente del servizio distribuisce il traffico in entrata a tutte le
JetStream HTTPrepliche. 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.Maxengineè un server JetStream che esegue l'inferenza con il batch continuo.
Prima di iniziare
- Accedi al tuo account Google Cloud. Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti gratuiti per l'esecuzione, il test e il deployment dei carichi di lavoro.
-
Nella pagina del selettore di progetti della console Google Cloud, seleziona o crea un progetto Google Cloud.
-
Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.
-
Attiva l'API richiesta.
-
Nella pagina del selettore di progetti della console Google Cloud, seleziona o crea un progetto Google Cloud.
-
Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.
-
Attiva l'API richiesta.
-
Assicurati di disporre dei seguenti ruoli nel progetto: roles/container.admin, roles/iam.serviceAccountAdmin
Verifica i ruoli
-
Nella console Google Cloud, vai alla pagina IAM.
Vai a IAM - Seleziona il progetto.
-
Nella colonna Entità, individua la riga contenente il tuo indirizzo email.
Se il tuo indirizzo email non è in questa colonna, significa che non disponi di alcun ruolo.
- Nella colonna Ruolo per la riga contenente il tuo indirizzo email, controlla se l'elenco dei ruoli include quelli richiesti.
Concedi i ruoli
-
Nella console Google Cloud, vai alla pagina IAM.
Vai a IAM - Seleziona il progetto.
- Fai clic su Concedi l'accesso.
- Nel campo Nuove entità, inserisci il tuo indirizzo email.
- Nell'elenco Seleziona un ruolo, scegli un ruolo.
- Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ciascun ruolo aggiuntivo.
- Fai clic su Salva.
-
- Assicurati di disporre di una quota sufficiente per otto chip TPU v5e PodSlice Lite. In questo tutorial utilizzerai le istanze on demand.
- Crea un account Kaggle, se non ne hai già uno.
Accedi al modello
Per ottenere l'accesso al modello Gemma per il deployment in GKE, devi prima firmare il contratto di licenza.
Firma il contratto di consenso per la licenza
Devi firmare il contratto di consenso per utilizzare Gemma. Segui queste istruzioni:
- Accedi alla pagina relativa al consenso del modello Gemma su Kaggle.com.
- Se non l'hai ancora fatto, accedi a Kaggle.
- Fai clic su Richiedi accesso.
- Nella sezione Scegli account per il consenso, seleziona Verifica tramite account Kaggle per utilizzare il tuo account Kaggle per il consenso.
- Accetta i Termini e condizioni del modello.
Generare un token di accesso
Per accedere al modello tramite Kaggle, è necessario un token API di Kaggle.
Se non ne hai già uno, segui questi passaggi per generare un nuovo token:
- Nel browser, vai alle impostazioni di Kaggle.
- Nella sezione API, fai clic su Crea nuovo token.
È stato scaricato un file denominato kaggle.json.
prepara l'ambiente
In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate su Google Cloud. Il software necessario per questo tutorial è preinstallato in Cloud Shell, tra cui kubectl e
gcloud CLI.
Per configurare il tuo ambiente con Cloud Shell, segui questi passaggi:
Nella console Google Cloud, avvia una sessione di Cloud Shell facendo clic su
Attiva Cloud Shell nella console Google Cloud. Viene avviata una sessione
nel riquadro inferiore della console Google Cloud.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 BUCKET_NAME=BUCKET_NAME export REGION=REGION export LOCATION=LOCATIONSostituisci i seguenti valori:
- PROJECT_ID: il tuo ID progetto Google Cloud.
- CLUSTER_NAME: il nome del tuo cluster GKE.
- BUCKET_NAME: il nome del bucket Cloud Storage.
Non è necessario specificare il prefisso
gs://. - 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 tipi di macchine TPU v5e (ad esempio
us-west1,us-west4,us-central1,us-east1,us-east5oeurope-west4). - (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.
crea e configura risorse Google Cloud
Segui queste istruzioni per creare le risorse richieste.
crea un cluster GKE
Puoi gestire 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
In Cloud Shell, esegui questo comando:
gcloud container clusters create-auto ${CLUSTER_NAME} \
--project=${PROJECT_ID} \
--region=${REGION} \
--release-channel=rapid \
--cluster-version=1.29
GKE crea un cluster Autopilot con nodi CPU e TPU, come richiesto dai carichi di lavoro di cui è stato eseguito il deployment.
Standard
Crea un cluster GKE Standard a livello di regione che utilizza la federazione delle identità per i carichi di lavoro per GKE.
gcloud container clusters create ${CLUSTER_NAME} \ --enable-ip-alias \ --machine-type=e2-standard-4 \ --num-nodes=2 \ --release-channel=rapid \ --workload-pool=${PROJECT_ID}.svc.id.goog \ --location=${REGION}La creazione del cluster potrebbe richiedere diversi minuti.
Esegui questo comando per creare un pool di nodi per il tuo cluster:
gcloud container node-pools create gemma-7b-tpu-nodepool \ --cluster=${CLUSTER_NAME} \ --machine-type=ct5lp-hightpu-8t \ --project=${PROJECT_ID} \ --num-nodes=2 \ --region=${REGION} \ --node-locations=${LOCATION}GKE crea un pool di nodi TPU v5e con una topologia
2x4e due nodi.
Crea un bucket Cloud Storage
In Cloud Shell, esegui questo comando:
gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}
Viene creato un bucket Cloud Storage per archiviare i file del modello scaricati da Kaggle.
Carica il token di accesso su Cloud Shell
In Cloud Shell, puoi caricare il token dell'API Kaggle nel tuo progetto Google Cloud:
- In Cloud Shell, fai clic su Altro > Carica.
- Seleziona File e fai clic su Scegli file.
- Apri il file
kaggle.json. - Fai clic su Carica.
Crea un secret di Kubernetes per le credenziali di Kaggle
In Cloud Shell, segui questi passaggi:
Configura
kubectlper comunicare con il cluster:gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}Crea un secret per archiviare le credenziali di Kaggle:
kubectl create secret generic kaggle-secret \ --from-file=kaggle.json
Configura l'accesso ai carichi di lavoro utilizzando la federazione di Workload Identity per GKE
Assegna un account di servizio Kubernetes all'applicazione e configuralo in modo che funzioni come account di servizio IAM.
Crea un account di servizio IAM per l'applicazione:
gcloud iam service-accounts create wi-jetstreamAggiungi un'associazione dei criteri 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.insightsCollectorServiceConsentire all'Account di servizio Kubernetes di assumere l'identità dell'account di servizio IAM aggiungendo un'associazione di criteri IAM tra i due account di servizio. Questa associazione consente a Kubernetes ServiceAccount 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]"Annota l'account di servizio Kubernetes con l'indirizzo email dell'account di servizio IAM:
kubectl annotate serviceaccount default \ iam.gke.io/gcp-service-account=wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com
Converti i checkpoint del modello
In questa sezione creerai un job per:
- Scarica il checkpoint di base Orbax da Kaggle.
- Carica il checkpoint in un bucket Cloud Storage.
- Converti il checkpoint in un checkpoint compatibile con MaxText.
- Annulla la scansione del checkpoint da utilizzare per la pubblicazione.
Esegui il deployment del job di conversione del checkpoint del modello
Segui queste istruzioni per scaricare e convertire i file checkpoint del modello Gemma 7B.
Crea il seguente manifest come
job-7b.yaml.apiVersion: batch/v1 kind: Job metadata: name: data-loader-7b spec: ttlSecondsAfterFinished: 30 template: spec: restartPolicy: Never containers: - name: inference-checkpoint image: us-docker.pkg.dev/cloud-tpu-images/inference/inference-checkpoint:v0.2.0 args: - -b=BUCKET_NAME - -m=google/gemma/maxtext/7b-it/2 volumeMounts: - mountPath: "/kaggle/" name: kaggle-credentials readOnly: true resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8 nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice volumes: - name: kaggle-credentials secret: defaultMode: 0400 secretName: kaggle-secretApplica il manifest:
kubectl apply -f job-7b.yamlVisualizza i log dal job:
kubectl logs -f jobs/data-loader-7bQuando il job viene completato, l'output è simile al seguente:
Successfully generated decode checkpoint at: gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items + echo -e '\nCompleted unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items' Completed unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
Esegui il deployment di JetStream
In questa sezione, eseguirai il deployment del container JetStream per gestire il modello Gemma.
Segui queste istruzioni per eseguire il deployment del modello ottimizzato per le istruzioni Gemma 7B.
Crea il seguente manifest
jetstream-gemma-deployment.yaml:apiVersion: apps/v1 kind: Deployment metadata: name: maxengine-server spec: replicas: 2 selector: matchLabels: app: maxengine-server template: metadata: labels: app: maxengine-server spec: nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice containers: - name: maxengine-server image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.0 args: - model_name=gemma-7b - tokenizer_path=assets/tokenizer.gemma - per_device_batch_size=4 - max_prefill_predict_length=1024 - max_target_length=2048 - async_checkpointing=false - ici_fsdp_parallelism=1 - ici_autoregressive_parallelism=-1 - ici_tensor_parallelism=1 - scan_layers=false - weight_dtype=bfloat16 - load_parameters_path=gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items ports: - containerPort: 9000 resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8 - name: jetstream-http image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.0 ports: - containerPort: 8000 --- apiVersion: v1 kind: Service metadata: name: jetstream-http-svc spec: selector: app: maxengine-server ports: - protocol: TCP port: 8000 targetPort: 8000Il file manifest imposta le seguenti proprietà chiave:
tokenizer_path: il percorso del tokenizzatore del modello.load_parameters_path: il percorso nel bucket Cloud Storage in cui sono archiviati i checkpoint.per_device_batch_size: la dimensione del batch di decodifica per dispositivo, dove un chip TPU equivale a un dispositivo.max_prefill_predict_length: la lunghezza massima della precompilazione quando si esegue l'autoregressione.max_target_length: la lunghezza massima della sequenza.model_name: il nome del modello (gemma-7b).ici_fsdp_parallelism: il numero di shard per il parallelismo dei dati completamente con sharding (FSDP).ici_tensor_parallelism: il numero di shard per il parallelismo tensore.ici_autoregressive_parallelism: il numero di shard per il parallelismo autoregressivo.scan_layers: scan layers boolean flag (boolean).weight_dtype: il tipo di dati sulla ponderazione (bfloat16).
Applica il manifest:
kubectl apply -f jetstream-gemma-deployment.yamlVerifica il deployment:
kubectl get deploymentL'output è simile al seguente:
NAME READY UP-TO-DATE AVAILABLE AGE maxengine-server 2/2 2 2 ##sVisualizza i log del server HTTP per verificare che il modello sia stato caricato e compilato. Il server potrebbe richiedere alcuni minuti per completare questa operazione.
kubectl logs deploy/maxengine-server -f -c jetstream-httpL'output è simile al seguente:
kubectl logs deploy/maxengine-server -f -c jetstream-http INFO: Started server process [1] INFO: Waiting for application startup. INFO: Application startup complete. INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)Visualizza i log di MaxEngine e verifica che la compilazione sia completata.
kubectl logs deploy/maxengine-server -f -c maxengine-serverL'output è simile al seguente:
2024-03-29 17:09:08,047 - jax._src.dispatch - DEBUG - Finished XLA compilation of jit(initialize) in 0.26236414909362793 sec 2024-03-29 17:09:08,150 - root - INFO - ---------Generate params 0 loaded.---------
Pubblica il modello
In questa sezione interagisci con il modello.
Configura il port forwarding
Puoi accedere al deployment 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 questi passaggi:
Per stabilire una sessione di port forwarding, esegui questo comando:
kubectl port-forward svc/jetstream-http-svc 8000:8000
Interagisci con il modello utilizzando curl
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": "\nfor data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its simplicity, readability, and extensive libraries for data wrangling, analysis, visualization, and machine learning.\n- Popular libraries include pandas, scikit-learn, and matplotlib.\n\n**2. R:**\n- Statistical programming language widely used for data analysis, visualization, and modeling.\n- Popular libraries include ggplot2, dplyr, and caret.\n\n**3. Java:**\n- Enterprise-grade language with strong performance and scalability.\n- Popular libraries include Spark, TensorFlow, and Weka.\n\n**4. C++:**\n- High-performance language often used for data analytics and machine learning models.\n- Popular libraries include TensorFlow, PyTorch, and OpenCV.\n\n**5. SQL:**\n- Relational database language essential for data wrangling and querying large datasets.\n- Popular tools" }
(Facoltativo) Interagisci con il modello tramite un'interfaccia della chat di Gradio
In questa sezione creerai un'applicazione di chat web che ti consente di interagire con il tuo modello ottimizzato per le istruzioni.
Gradio è una libreria Python con un wrapper
ChatInterface che crea le interfacce utente per i chatbot.
Eseguire il deployment dell'interfaccia della chat
In Cloud Shell, salva il manifest seguente come
gradio.yaml:apiVersion: apps/v1 kind: Deployment metadata: name: gradio labels: app: gradio spec: replicas: 1 selector: matchLabels: app: gradio template: metadata: labels: app: gradio spec: containers: - name: gradio image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.3 resources: requests: cpu: "512m" memory: "512Mi" limits: cpu: "1" memory: "512Mi" env: - name: CONTEXT_PATH value: "/generate" - name: HOST value: "http://jetstream-http-svc:8000" - name: LLM_ENGINE value: "max" - name: MODEL_ID value: "gemma" - name: USER_PROMPT value: "<start_of_turn>user\nprompt<end_of_turn>\n" - name: SYSTEM_PROMPT value: "<start_of_turn>model\nprompt<end_of_turn>\n" ports: - containerPort: 7860 --- apiVersion: v1 kind: Service metadata: name: gradio spec: selector: app: gradio ports: - protocol: TCP port: 8080 targetPort: 7860 type: ClusterIPApplica il manifest:
kubectl apply -f gradio.yamlAttendi che il deployment sia disponibile:
kubectl wait --for=condition=Available --timeout=300s deployment/gradio
Utilizzare l'interfaccia della chat
In Cloud Shell, esegui questo comando:
kubectl port-forward service/gradio 8080:8080Viene creato un port forwarding da Cloud Shell al servizio Gradio.
Fai clic sul pulsante
Anteprima web in alto a destra della barra delle applicazioni di Cloud Shell. Fai clic su Anteprima sulla porta 8080. Nel browser si apre una nuova scheda.Interagisci con Gemma utilizzando l'interfaccia della chat Gradio. Aggiungi un prompt e fai clic su Invia.
Risoluzione dei problemi
- Se ricevi il messaggio
Empty reply from server, è possibile che il container non abbia completato il download dei dati del modello. Controlla di nuovo nei log del pod il messaggioConnectedche indica che il modello è pronto per essere pubblicato. - Se vedi
Connection refused, verifica che il port forwarding 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 questi 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
gcloud storage rm --recursive gs://BUCKET_NAME
Passaggi successivi
- Scopri come eseguire modelli Gemma su GKE e come eseguire carichi di lavoro AI/ML ottimizzati con le funzionalità di orchestrazione della piattaforma GKE.
- Scopri di più sulle TPU in GKE.
- Esplora il repository di GitHub JetStream.
- Esplora Vertex AI Model Garden.