Gestisci modelli aperti Gemma utilizzando TPU su GKE con Saxml

Standard

---

Questa guida mostra come gestire un modello linguistico di grandi dimensioni (LLM) di Gemma utilizzando TPU (Tensor Processing Unit) su Google Kubernetes Engine (GKE) con Saxml. In questa guida scaricherai i modelli Gemma ottimizzati dell'istruzione con parametri 2 miliardi e 7 miliardi in Cloud Storage e ne eseguirai il deployment su un cluster GKE Standard utilizzando container che eseguono Saxml.

Questa guida è un buon punto di partenza se hai bisogno delle funzionalità di scalabilità, resilienza ed economicità offerte da Kubernetes quando esegui il deployment del tuo modello su Saxml. Se hai bisogno di una piattaforma IA gestita e unificata per creare e pubblicare rapidamente modelli ML a costi contenuti, ti consigliamo di provare la nostra soluzione di deployment Vertex AI.

Contesto

Utilizzando Gemma mediante TPU su GKE con Saxml, puoi implementare una solida soluzione di gestione di inferenza pronta per la produzione con tutti i vantaggi di Kubernetes gestito, tra cui scalabilità efficiente e disponibilità superiore. In questa sezione vengono descritte le tecnologie principali utilizzate in questo tutorial.

Gemma

Gemma è un insieme di modelli di AI generativa Leggeri e disponibili, disponibili con licenza aperta. Questi modelli di IA sono disponibili per l'esecuzione nelle tue applicazioni, nel tuo hardware, nei tuoi dispositivi mobili o nei tuoi 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 i framework di elaborazione dei dati 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à attuale delle versioni di TPU con l'architettura di sistema di Cloud TPU.
2. Scopri di più sulle TPU in GKE.

Questo tutorial serve i modelli Gemma 2B e Gemma 7B. GKE ospita questi modelli sui seguenti pool di nodi TPU v5e con host singolo:

• Gemma 2B: modello ottimizzato per le istruzioni ospitato in un pool di nodi TPU v5e con topologia 1x1 che rappresenta un chip TPU. Il tipo di macchina per i nodi è ct5lp-hightpu-1t.
• Gemma 7B: modello ottimizzato per le istruzioni ospitato in un pool di nodi TPU v5e con topologia 2x2 che rappresenta quattro chip TPU. Il tipo di macchina per i nodi è ct5lp-hightpu-4t.

Saxml

Saxml è un sistema sperimentale che gestisce i modelli Paxml, JAX e PyTorch per l'inferenza. Il sistema Saxml comprende i seguenti componenti:

• Cella Saxml o cluster Sax: comprende un server di amministrazione e un gruppo di server modello. Il server di amministrazione tiene traccia dei server dei modelli, assegna i modelli pubblicati ai server dei modelli da utilizzare e aiuta i client a individuare i server dei modelli che pubblicano modelli specifici.
• Client Saxml: l'interfaccia di programmazione rivolta all'utente per il sistema Saxml. Il client Saxml include uno strumento a riga di comando (saxutil) e una suite di librerie client in Python, C++ e Go.

In questo tutorial utilizzerai anche il server HTTP Saxml. Saxml HTTP Server è un server HTTP personalizzato che incapsula la libreria client Python Saxml ed espone le API REST per l'interazione con il sistema Saxml, inclusi endpoint per pubblicare, elencare, annullare la pubblicazione di modelli e generare previsioni.

Obiettivi

Questo tutorial è destinato ai clienti di IA generativa che utilizzano JAX, agli utenti nuovi o esistenti di GKE, agli ML engineer, agli ingegneri MLOps (DevOps) o agli amministratori di piattaforma interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per gestire Gemma.

Questo tutorial illustra i seguenti passaggi:

1. Prepara un cluster GKE Standard con la topologia TPU consigliata in base alle caratteristiche del modello.
2. Eseguire il deployment dei componenti Saxml su GKE.
3. Ottieni e pubblica il modello di parametri Gemma 2B o Gemma 7B.
4. Pubblica e interagisci con i modelli pubblicati.

Architettura

Questa sezione descrive l'architettura GKE utilizzata in questo tutorial. L'architettura comprende un cluster GKE Standard che esegue il provisioning delle TPU e ospita i componenti Saxml per il deployment e la gestione di modelli Gemma 2B o 7B. Il seguente diagramma mostra i componenti di questa architettura:

Questa architettura include i seguenti componenti:

• Un cluster GKE Standard a livello di zona.
• Un pool di nodi della sezione TPU con singolo host che dipende dal modello Gemma che vuoi gestire:
  • Gemma 2B: configurata con una TPU v5e con una topologia 1x1. Un'istanza del server del modello Saxml è configurata per l'utilizzo di questo pool di nodi.
  • Gemma 7B: configurata con una TPU v5e con una topologia 2x2. Un'istanza del server del modello Saxml è configurata per l'utilizzo di questo pool di nodi.
• Un pool di nodi CPU predefinito in cui vengono distribuiti il server di amministrazione Saxml e il server HTTP Saxml.
• Due bucket Cloud Storage:
  • Un bucket Cloud Storage archivia lo stato gestito da un server di amministrazione.
  • Un bucket Cloud Storage archivia i checkpoint del modello Gemma.

Questa architettura presenta le seguenti caratteristiche:

• Un Artifact Registry pubblico gestisce le immagini container per i componenti Saxml.
• Il cluster GKE utilizza la Federazione delle identità per i carichi di lavoro per GKE. Tutti i componenti Saxml utilizzano una federazione delle identità per i carichi di lavoro che integra un account di servizio IAM per accedere a servizi esterni come i bucket Cloud Storage.
• I log generati dai componenti Saxml sono integrati in Cloud Logging.
• Puoi utilizzare Cloud Monitoring per analizzare le metriche delle prestazioni dei pool di nodi GKE creati da questo tutorial.

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.

  Nota: se non prevedi di conservare le risorse create in questa procedura, crea un progetto anziché selezionarne uno esistente. Dopo aver completato questi passaggi, puoi eliminare il progetto, rimuovendo tutte le risorse associate.

  Vai al selettore progetti

• Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

• Attiva l'API richiesta.

  Abilita l'API

• Nella pagina del selettore di progetti della console Google Cloud, seleziona o crea un progetto Google Cloud.

  Nota: se non prevedi di conservare le risorse create in questa procedura, crea un progetto anziché selezionarne uno esistente. Dopo aver completato questi passaggi, puoi eliminare il progetto, rimuovendo tutte le risorse associate.

  Vai al selettore progetti

• Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

• Attiva l'API richiesta.

  Abilita l'API

• Assicurati di disporre dei seguenti ruoli nel progetto: roles/container.admin, roles/iam.serviceAccountAdmin

  Verifica i ruoli

  1. Nella console Google Cloud, vai alla pagina IAM.

     Vai a IAM
  2. Seleziona il progetto.

  3. 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.

  4. Nella colonna Ruolo per la riga contenente il tuo indirizzo email, controlla se l'elenco dei ruoli include quelli richiesti.

  Concedi i ruoli

  1. Nella console Google Cloud, vai alla pagina IAM.

     Vai a IAM
  2. Seleziona il progetto.
  3. Fai clic su Concedi l'accesso.
  4. Nel campo Nuove entità, inserisci il tuo indirizzo email.
  5. Nell'elenco Seleziona un ruolo, scegli un ruolo.
  6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ciascun ruolo aggiuntivo.
  7. Fai clic su Salva.

• Assicurati di disporre di quote sufficienti per 5 chip TPU v5e. In questo tutorial utilizzerai le istanze on demand.
• Crea un account Kaggle, se non ne hai già uno.

prepara l'ambiente per Gemma

Avvia Cloud Shell

In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate su Google Cloud. Cloud Shell è preinstallato con il software necessario per questo tutorial, inclusi kubectl e gcloud CLI.

1. Nella console Google Cloud, avvia un'istanza di Cloud Shell:
   Apri Cloud Shell

2. Imposta le variabili di ambiente predefinite:

   gcloud config set project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   export LOCATION=LOCATION
   export CLUSTER_NAME=saxml-tpu
   

   Sostituisci i seguenti valori:

   • PROJECT_ID: il tuo ID progetto Google Cloud.
   • LOCATION: il nome della zona di Compute Engine in cui sono disponibili i tipi di macchina TPU v5e.

Crea un cluster GKE Standard

In questa sezione creerai il cluster GKE e il pool di nodi.

Gemma 2Bit

Utilizza Cloud Shell per:

1. Crea un cluster standard 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=${LOCATION}
   

   La creazione del cluster potrebbe richiedere diversi minuti.

2. Crea un pool di nodi TPU v5e con una topologia 1x1 e un nodo:

   gcloud container node-pools create tpu-v5e-1x1 \
       --cluster=${CLUSTER_NAME} \
       --machine-type=ct5lp-hightpu-1t \
       --num-nodes=1 \
       --location=${LOCATION}
   

   Gestisci il modello Gemma 2B in questo pool di nodi.

Gemma 7B-it

Utilizza Cloud Shell per:

1. Crea un cluster standard 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=${LOCATION}
   

   La creazione del cluster potrebbe richiedere diversi minuti.

2. Crea un pool di nodi TPU v5e con una topologia 2x2 e un nodo:

   gcloud container node-pools create tpu-v5e-2x2 \
       --cluster=${CLUSTER_NAME} \
       --machine-type=ct5lp-hightpu-4t \
       --num-nodes=1 \
       --location=${LOCATION}
   

   Gestisci il modello Gemma 7B in questo pool di nodi.

crea i bucket Cloud Storage

Crea due bucket Cloud Storage per gestire lo stato del server di amministrazione Saxml e i checkpoint del modello.

In Cloud Shell, esegui questo comando:

1. Creare un bucket Cloud Storage per archiviare le configurazioni del server di amministrazione Saxml.

   gcloud storage buckets create gs://ADMIN_BUCKET_NAME
   

   Sostituisci ADMIN_BUCKET_NAME con il nome del bucket Cloud Storage in cui è archiviato il server di amministrazione Saxml.

2. Crea un bucket Cloud Storage per archiviare i checkpoint del modello:

   gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
   

   Sostituisci CHECKPOINTS_BUCKET_NAME con il nome del bucket Cloud Storage in cui sono archiviati i checkpoint del modello.

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.

1. Configura kubectl per comunicare con il cluster:

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

2. Crea un ServiceAccount Kubernetes per l'applicazione da utilizzare:

   gcloud iam service-accounts create wi-sax
   

3. Aggiungi un'associazione dei criteri IAM per il tuo account di servizio IAM per leggere e scrivere in Cloud Storage:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
       --role roles/storage.objectUser
   
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
       --role roles/storage.insightsCollectorService
   

4. Consentire 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-sax@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/iam.workloadIdentityUser \
       --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
   

5. 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-sax@${PROJECT_ID}.iam.gserviceaccount.com
   

Accedi al modello

Per ottenere l'accesso ai modelli Gemma per il deployment su GKE, devi accedere alla piattaforma Kaggle, firmare il contratto di consenso di licenza e ottenere un token API Kaggle. In questo tutorial utilizzerai un secret di Kubernetes per le credenziali di Kaggle.

Firma il contratto di consenso per la licenza

Devi firmare il contratto di consenso per utilizzare Gemma. Segui queste istruzioni:

1. Accedi alla pagina relativa al consenso del modello su Kaggle.com.
2. Se non l'hai ancora fatto, accedi a Kaggle.
3. Fai clic su Richiedi accesso.
4. Nella sezione Scegli account per il consenso, seleziona Verifica tramite account Kaggle per utilizzare il tuo account Kaggle per il consenso.
5. Accetta i Termini e condizioni del modello.

Generare un token di accesso

Per accedere al modello tramite Kaggle, è necessario un token API Kaggle.

Se non ne hai già uno, segui questi passaggi per generare un nuovo token:

1. Nel browser, vai alle impostazioni di Kaggle.
2. Nella sezione API, fai clic su Crea nuovo token.

È stato scaricato un file denominato kaggle.json.

Carica il token di accesso su Cloud Shell

In Cloud Shell, puoi caricare il token dell'API Kaggle nel tuo progetto Google Cloud:

1. In Cloud Shell, fai clic su Altro > Carica.
2. Seleziona File e fai clic su Scegli file.
3. Apri il file kaggle.json.
4. Fai clic su Carica.

Crea il secret di Kubernetes per le credenziali di Kaggle

In Cloud Shell, segui questi passaggi:

1. Configura kubectl per comunicare con il cluster:

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

2. Crea un secret per archiviare le credenziali di Kaggle:

   kubectl create secret generic kaggle-secret \
       --from-file=kaggle.json
   

Deployment di Saxml

In questa sezione, eseguirai il deployment del server di amministrazione Saxml, dei server del modello e del server HTTP.

Deployment del server di amministrazione Saxml

1. Crea il seguente manifest saxml-admin-server.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: sax-admin-server
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: sax-admin-server
     template:
       metadata:
         labels:
           app: sax-admin-server
       spec:
         hostNetwork: false
         containers:
         - name: sax-admin-server
           image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.2.0
           securityContext:
             privileged: true
           ports:
           - containerPort: 10000
           env:
           - name: GSBUCKET
             value: ADMIN_BUCKET_NAME

   Sostituisci ADMIN_BUCKET_NAME con il nome del bucket creato nella sezione Crea bucket Cloud Storage. Non includere il prefisso gs://.

2. Applica il manifest:

   kubectl apply -f saxml-admin-server.yaml
   

3. Verifica il deployment del server di amministrazione:

   kubectl get deployment
   

   L'output dovrebbe essere simile al seguente:

   NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
   sax-admin-server                  1/1     1            1           ##s
   

Deployment del server del modello Saxml

Segui queste istruzioni per eseguire il deployment del server del modello per il modello Gemma 2B o Gemma 7B.

Gemma 2Bit

1. Crea il seguente manifest saxml-model-server-1x1.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: sax-model-server-v5e-1x1
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: gemma-server
     strategy:
       type: Recreate
     template:
       metadata:
         labels:
           app: gemma-server
           ai.gke.io/model: gemma-2b-it
           ai.gke.io/inference-server: saxml
           examples.ai.gke.io/source: user-guide
       spec:
         nodeSelector:
           cloud.google.com/gke-tpu-topology: 1x1
           cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
         hostNetwork: false
         restartPolicy: Always
         containers:
         - name: inference-server
           image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
           args:
           - "--jax_platforms=tpu"
           - "--platform_chip=tpuv5e"
           - "--platform_topology=1x1"
           - "--port=10001"
           - "--sax_cell=/sax/test"
           ports:
           - containerPort: 10001
           securityContext:
             privileged: true
           env:
           - name: SAX_ROOT
             value: "gs://ADMIN_BUCKET_NAME/sax-root"
           resources:
             requests:
               google.com/tpu: 1
             limits:
               google.com/tpu: 1

   Sostituisci ADMIN_BUCKET_NAME con il nome del bucket creato nella sezione Crea bucket Cloud Storage. Non includere il prefisso gs://.

2. Applica il manifest:

   kubectl apply -f saxml-model-server-1x1.yaml
   

3. Verifica lo stato del deployment del server del modello:

   kubectl get deployment
   

   L'output dovrebbe essere simile al seguente:

   NAME                                              READY   STATUS    RESTARTS   AGE
   sax-admin-server                                  1/1     Running   0          ##m
   sax-model-server-v5e-1x1                          1/1     Running   0          ##s
   

Operazione riuscita: hai eseguito il deployment di un server del modello su un pool di nodi TPU v5e con una sezione 1 x 1 per gestire il modello Gemma 2B.

Gemma 7B-it

1. Crea il seguente manifest saxml-model-server-2x2.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: sax-model-server-v5e-2x2
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: gemma-server
     strategy:
       type: Recreate
     template:
       metadata:
         labels:
           app: gemma-server
           ai.gke.io/model: gemma-7b-it
           ai.gke.io/inference-server: saxml
           examples.ai.gke.io/source: user-guide
       spec:
         nodeSelector:
           cloud.google.com/gke-tpu-topology: 2x2
           cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
         hostNetwork: false
         restartPolicy: Always
         containers:
         - name: inference-server
           image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
           args:
           - "--jax_platforms=tpu"
           - "--platform_chip=tpuv5e"
           - "--platform_topology=2x2"
           - "--port=10001"
           - "--sax_cell=/sax/test"
           ports:
           - containerPort: 10001
           securityContext:
             privileged: true
           env:
           - name: SAX_ROOT
             value: "gs://ADMIN_BUCKET_NAME/sax-root"
           resources:
             requests:
               google.com/tpu: 4
             limits:
               google.com/tpu: 4

   Sostituisci ADMIN_BUCKET_NAME con il nome del bucket creato nella sezione Crea bucket Cloud Storage. Non includere il prefisso gs://.

2. Applica il manifest:

   kubectl apply -f saxml-model-server-2x2.yaml
   

3. Verifica lo stato del deployment del server del modello:

   kubectl get deployment
   

   L'output dovrebbe essere simile al seguente:

   NAME                                              READY   STATUS    RESTARTS   AGE
   sax-admin-server                                  1/1     Running   0          ##m
   sax-model-server-v5e-2x2                          1/1     Running   0          ##s
   

Operazione riuscita: hai eseguito il deployment di un server del modello su un pool di nodi TPU v5e con una sezione 2x2 per gestire il modello Gemma 7B.

Deployment del server HTTP Saxml

In questa sezione, eseguirai il deployment del server HTTP Saxml e creerai un servizio IP cluster che utilizzi per accedere al server.

1. Crea il seguente manifest saxml-http.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: sax-http
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: sax-http
     template:
       metadata:
         labels:
           app: sax-http
       spec:
         hostNetwork: false
         containers:
         - name: sax-http
           image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.2.0
           imagePullPolicy: Always
           ports:
           - containerPort: 8888
           env:
           - name: SAX_ROOT
             value: "gs://ADMIN_BUCKET_NAME/sax-root"
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: sax-http-svc
   spec:
     selector:
       app: sax-http
     ports:
     - protocol: TCP
       port: 8888
       targetPort: 8888
     type: ClusterIP

   Sostituisci ADMIN_BUCKET_NAME con il nome del bucket Cloud Storage in cui è archiviato il server di amministrazione Saxml.

2. Applica il manifest:

   kubectl apply -f saxml-http.yaml
   

3. Verifica lo stato del deployment del server HTTP Saxml:

   kubectl get deployment
   

Gemma 2Bit

L'output dovrebbe essere simile al seguente:

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-1x1                          1/1     Running   0          ##m
sax-http                                          1/1     Running   0          ##s

Gemma 7B-it

L'output dovrebbe essere simile al seguente:

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-2x2                          1/1     Running   0          ##m
sax-http                                          1/1     Running   0          ##s

Scarica il checkpoint del modello

In questa sezione eseguirai un job Kubernetes che recupera, scarica e archivia il checkpoint del modello. Segui i passaggi in base al modello Gemma che vuoi utilizzare:

Gemma 2Bit

1. Crea il seguente manifest job-2b.yaml:

   apiVersion: v1
   kind: ConfigMap
   metadata:
     name: fetch-model-scripts
   data:
     fetch_model.sh: |-
       #!/usr/bin/bash -x
       pip install kaggle --break-system-packages && \
   
       MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
       VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
   
       mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
       kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
       echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
   
       gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
       echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
   ---
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: data-loader-2b
     labels:
       app: data-loader-2b
   spec:
     ttlSecondsAfterFinished: 120
     template:
       metadata:
         labels:
           app: data-loader-2b
       spec:
         restartPolicy: OnFailure
         containers:
         - name: gcloud
           image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
           command:
           - /scripts/fetch_model.sh
           env:
           - name: BUCKET_NAME
             value: CHECKPOINTS_BUCKET_NAME
           - name: KAGGLE_CONFIG_DIR
             value: /kaggle
           - name: MODEL_PATH
             value: "google/gemma/pax/2b-it/2"
           volumeMounts:
           - mountPath: "/kaggle/"
             name: kaggle-credentials
             readOnly: true
           - mountPath: "/scripts/"
             name: scripts-volume
             readOnly: true
         volumes:
         - name: kaggle-credentials
           secret:
             defaultMode: 0400
             secretName: kaggle-secret
         - name: scripts-volume
           configMap:
             defaultMode: 0700
             name: fetch-model-scripts

   Sostituisci CHECKPOINTS_BUCKET_NAME con il nome del bucket creato nella sezione Crea bucket Cloud Storage. Non includere il prefisso gs://.

2. Applica il manifest:

   kubectl apply -f job-2b.yaml
   

3. Attendi il completamento del job:

   kubectl wait --for=condition=complete --timeout=180s job/data-loader-2b
   

   L'output è simile al seguente:

   job.batch/data-loader-2b condition met
   

4. Verifica che il job sia stato completato correttamente:

   kubectl get job/data-loader-2b
   

   L'output è simile al seguente:

   NAME             COMPLETIONS   DURATION   AGE
   data-loader-2b   1/1           ##s        #m##s
   

5. Visualizza i log per il job:

   kubectl logs --follow job/data-loader-2b
   

I checkpoint sono stati caricati su gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

Gemma 7B-it

1. Crea il seguente manifest job-7b.yaml:

   apiVersion: v1
   kind: ConfigMap
   metadata:
     name: fetch-model-scripts
   data:
     fetch_model.sh: |-
       #!/usr/bin/bash -x
       pip install kaggle --break-system-packages && \
   
       MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
       VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
   
       mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
       kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
       echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
   
       gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
       echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
   ---
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: data-loader-7b
     labels:
       app: data-loader-7b
   spec:
     ttlSecondsAfterFinished: 120
     template:
       metadata:
         labels:
           app: data-loader-7b
       spec:
         restartPolicy: OnFailure
         containers:
         - name: gcloud
           image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
           command:
           - /scripts/fetch_model.sh
           env:
           - name: BUCKET_NAME
             value: CHECKPOINTS_BUCKET_NAME
           - name: KAGGLE_CONFIG_DIR
             value: /kaggle
           - name: MODEL_PATH
             value: "google/gemma/pax/7b-it/2"
           volumeMounts:
           - mountPath: "/kaggle/"
             name: kaggle-credentials
             readOnly: true
           - mountPath: "/scripts/"
             name: scripts-volume
             readOnly: true
         volumes:
         - name: kaggle-credentials
           secret:
             defaultMode: 0400
             secretName: kaggle-secret
         - name: scripts-volume
           configMap:
             defaultMode: 0700
             name: fetch-model-scripts

   Sostituisci CHECKPOINTS_BUCKET_NAME con il nome del bucket creato nella sezione Crea bucket Cloud Storage. Includi il prefisso gs://.

2. Applica il manifest:

   kubectl apply -f job-7b.yaml
   

3. Attendi il completamento del job:

   kubectl wait --for=condition=complete --timeout=360s job/data-loader-7b
   

   L'output è simile al seguente:

   job.batch/data-loader-7b condition met
   

4. Verifica che il job sia stato completato correttamente:

   kubectl get job/data-loader-7b
   

   L'output è simile al seguente:

   NAME             COMPLETIONS   DURATION   AGE
   data-loader-7b   1/1           ##s        #m##s
   

5. Visualizza i log per il job:

   kubectl logs --follow job/data-loader-7b
   

I checkpoint sono stati caricati su gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

Esporre il server HTTP Saxml

Puoi accedere al server HTTP Saxml tramite un 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:

1. Stabilisci una sessione di port forwarding:

   kubectl port-forward service/sax-http-svc 8888:8888
   

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

   curl -s localhost:8888
   

   L'output è simile al seguente:

   {
       "Message": "HTTP Server for SAX Client"
   }
   

Il server HTTP Saxml incapsula l'interfaccia client nel sistema Saxml e la espone tramite un set di API REST. Userai queste API per pubblicare, gestire e interfacciarti con i modelli Gemma 2B e Gemma 7B.

Pubblica il modello Gemma

Dopo aver configurato i componenti Saxml, puoi pubblicare il modello Gemma su un server modello eseguito in un pool di nodi della sezione TPU. Utilizzi l'API publish del server HTTP Saxml per pubblicare un modello. Per ulteriori informazioni sull'API del server HTTP Saxml, consulta API HTTP Saxml. Per pubblicare il modello a parametri Gemma 2B o 7B, segui questi passaggi.

Gemma 2Bit

1. Assicurati che la sessione di port forwarding sia ancora attiva:

   curl -s localhost:8888
   

2. Pubblica il parametro Gemma 2B:

   curl --request POST \
   --header "Content-type: application/json" \
   -s \
   localhost:8888/publish \
   --data \
   '{
       "model": "/sax/test/gemma2bfp16",
       "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
       "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
       "replicas": "1"
   }'
   

   L'output è simile al seguente:

   {
       "model": "/sax/test/gemma2bfp16",
       "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
       "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
       "replicas": 1
   }
   

   Il deployment del modello può richiedere alcuni minuti.

3. Monitora l'avanzamento osservando i log in un pod di server del modello del deployment sax-model-server-v5e-1x1.

   kubectl logs --follow deployment/sax-model-server-v5e-1x1
   

   Il completamento di questo deployment può richiedere fino a cinque minuti. Attendi finché non viene visualizzato un messaggio simile al seguente:

   I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
   I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma2bfp16
   

4. Verifica di poter accedere al modello visualizzando le relative informazioni:

   curl --request GET \
   --header "Content-type: application/json" \
   -s \
   localhost:8888/listcell \
   --data \
   '{
       "model": "/sax/test/gemma2bfp16"
   }'
   

   L'output è simile al seguente:

   {
       "model": "/sax/test/gemma2bfp16",
       "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
       "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
       "max_replicas": 1,
       "active_replicas": 1
   }
   

Successo: hai pubblicato un modello Gemma 2B. Ora puoi interagire con il modello.

Gemma 7B-it

1. Assicurati che la sessione di port forwarding sia ancora attiva:

   curl -s localhost:8888
   

2. Pubblica il parametro Gemma 7B:

   curl --request POST \
   --header "Content-type: application/json" \
   -s \
   localhost:8888/publish \
   --data \
   '{
       "model": "/sax/test/gemma7bfp16",
       "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
       "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
       "replicas": "1"
   }'
   

   L'output è simile al seguente:

   {
       "model": "/sax/test/gemma7bfp16",
       "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
       "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
       "replicas": 1
   }
   

   Il deployment del modello può richiedere alcuni minuti.

3. Monitora l'avanzamento osservando i log in un pod di server del modello del deployment sax-model-server-v5e-2x2.

   kubectl logs --follow deployment/sax-model-server-v5e-2x2
   

   Attendi finché non viene visualizzato un messaggio simile al seguente:

   I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
   I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma7bfp16
   

4. Verifica che il modello sia stato pubblicato visualizzando le relative informazioni:

   curl --request GET \
   --header "Content-type: application/json" \
   -s \
   localhost:8888/listcell \
   --data \
   '{
       "model": "/sax/test/gemma7bfp16"
   }'
   

   L'output è simile al seguente:

   {
       "model": "/sax/test/gemma7bfp16",
       "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
       "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
       "max_replicas": 1,
       "active_replicas": 1
   }
   

Successo: hai pubblicato un modello Gemma 7B. Ora puoi interagire con il modello.

Utilizza il modello

Puoi interagire con i modelli Gemma 2B o 7B. Utilizza l'API generate del server HTTP Saxml per inviare un prompt al modello.

Gemma 2Bit

Per fornire una richiesta di prompt, utilizza l'endpoint generate del server HTTP Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma2bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

Di seguito è riportato un esempio di risposta del modello. L'output può variare in ogni prompt:

[
    [
        "\n\n1. **Python**\n2. **JavaScript**\n3. **Java**\n4. **C++**\n5. **Go**",
        -3.0704939365386963
    ]
]

Puoi eseguire il comando con diversi parametri query. Puoi anche modificare parametri aggiuntivi come temperature, top_k e topc_p utilizzando l'API generate. Per scoprire di più sull'API del server HTTP Saxml, consulta API HTTP Saxml.

Gemma 7B-it

Per fornire una richiesta di prompt, utilizza l'endpoint generate del server HTTP Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma7bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

Di seguito è riportato un esempio di risposta del modello. L'output può variare in ogni prompt:

[
    [
        "\n\n**1. JavaScript**\n\n* Most widely used language on the web.\n* Used for front-end development, such as websites and mobile apps.\n* Extensive libraries and frameworks available.\n\n**2. Python**\n\n* Known for its simplicity and readability.\n* Versatile, used for various tasks, including data science, machine learning, and web development.\n* Large and active community.\n\n**3. Java**\n\n* Object-oriented language widely used in enterprise applications.\n* Used for web applications, mobile apps, and enterprise software.\n* Strong ecosystem and support.\n\n**4. Go**\n\n",
        -16.806324005126953
    ]
]

Puoi eseguire il comando con diversi parametri query. Puoi anche modificare parametri aggiuntivi come temperature, top_k e topc_p utilizzando l'API generate. Per scoprire di più sull'API del server HTTP Saxml, consulta API HTTP Saxml.

Annulla la pubblicazione del modello

Per annullare la pubblicazione del modello, procedi nel seguente modo:

Gemma 2Bit

Per annullare la pubblicazione del modello Gemma 2B-it, esegui questo comando:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

L'output potrebbe essere simile al seguente:

{
    "model": "/sax/test/gemma2bfp16"
}

Puoi eseguire il comando con diversi prompt passati nel parametro query.

Gemma 7B-it

Per annullare la pubblicazione del modello Gemma 7B-it, esegui questo comando:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

L'output potrebbe essere simile al seguente:

{
    "model": "/sax/test/gemma7bfp16"
}

Puoi eseguire il comando con diversi prompt passati nel parametro query.

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 messaggio Connected che indica che il modello è pronto per la pubblicazione.
• 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 questo comando:

gcloud container clusters delete ${CLUSTER_NAME} --location=${LOCATION}
gcloud iam service-accounts delete --quiet wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
gcloud storage rm --recursive gs://ADMIN_BUCKET_NAME
gcloud storage rm --recursive gs://CHECKPOINTS_BUCKET_NAME

Sostituisci quanto segue:

• ADMIN_BUCKET_NAME: il nome del bucket Cloud Storage in cui è archiviato il server di amministrazione Saxml.
• CHECKPOINTS_BUCKET_NAME: il nome del bucket Cloud Storage in cui sono archiviati i checkpoint del modello.

Passaggi successivi

• Scopri di più sulle TPU in GKE.
• Esplora il repository GitHub di Saxml.
• Esplora Vertex AI Model Garden.
• Scopri come eseguire carichi di lavoro AI/ML ottimizzati con le funzionalità di orchestrazione della piattaforma GKE.
• Esplora le architetture di riferimento, i diagrammi e le best practice su Google Cloud. Visita il nostro Cloud Architecture Center.