Pubblicare modelli aperti Gemma utilizzando TPU su Vertex AI Prediction con Saxml

Anteprima

Questa funzionalità è soggetta ai "Termini delle Offerte pre-GA" nella sezione Termini generali del servizio dei Termini specifici dei servizi. Le funzionalità pre-GA sono disponibili "così come sono" e potrebbero avere un supporto limitato. Per saperne di più, consulta le descrizioni della fase di lancio.

Questa guida mostra come pubblicare un modello LLM (Large Language Model) di Gemma open models utilizzando le TPU (Tensor Processing Unit) in Vertex AI Prediction con Saxml. In questa guida, scarichi i modelli Gemma ottimizzati per le istruzioni sui parametri 2B e 7B in Cloud Storage ed esegui il loro deployment in Vertex AI Prediction, che esegue Saxml su TPU.

Sfondo

Pubblicando Gemma utilizzando TPU su Vertex AI Prediction con Saxml. Puoi usufruire di una soluzione di AI gestita che si occupa dell'infrastruttura di basso livello e offre un modo economico per pubblicare LLM. Questa sezione descrive le tecnologie chiave utilizzate in questo tutorial.

Gemma

Gemma è un insieme di modelli di AI (intelligenza artificiale) generativi, leggeri e disponibili pubblicamente rilasciati con una licenza aperta. Questi modelli di IA possono essere eseguiti nelle tue applicazioni, nel tuo hardware, sui 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.

Saxml

Saxml è un sistema sperimentale che serve modelli Paxml, JAX e PyTorch per l'inferenza. Ai fini di questo tutorial, spiegheremo come pubblicare Gemma su TPU più convenienti per Saxml. La configurazione per le GPU è simile. Saxml offre script per creare container per Vertex AI Prediction che utilizzeremo in questo tutorial.

TPU

Le TPU sono circuiti integrati per applicazioni specifiche (ASIC) sviluppati da Google e utilizzati per accelerare framework di elaborazione dei dati come TensorFlow, PyTorch e JAX.

Questo tutorial è rivolto ai modelli Gemma 2B e Gemma 7B. Vertex AI Prediction ospita questi modelli sui seguenti pool di nodi TPU v5e a host singolo:

• Gemma 2B: 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: 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.

Prima di iniziare

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

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. Enable the Vertex AI API.

   Enable the API

5. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

6. Make sure that billing is enabled for your Google Cloud project.

7. Enable the Vertex AI API.

   Enable the API

8. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

Questo tutorial presuppone che tu stia utilizzando Cloud Shell per interagire con Google Cloud. Se vuoi utilizzare una shell diversa da Cloud Shell, esegui la seguente configurazione aggiuntiva:

1. Install the Google Cloud CLI.

2. To initialize the gcloud CLI, run the following command:

   gcloud init

3. Assicurati di disporre di una quota sufficiente per i chip TPU v5e per Vertex AI Prediction. Per impostazione predefinita, questa quota è 0. Per una topologia 1x1, deve essere 1. Per 2x2, deve essere 4. Per eseguire entrambe le topologie, deve essere 5.
4. Se non ne hai già uno, crea un account Kaggle.

Ottieni l'accesso al modello

Tieni presente che Cloud Shell potrebbe non avere risorse sufficienti per scaricare i pesi del modello. In questo caso, puoi creare un'istanza di Vertex AI Workbench per eseguire l'attività.

Per accedere ai modelli Gemma per il deployment in Vertex AI Prediction, devi accedere alla piattaforma Kaggle, firmare il contratto di consenso alla licenza e ottenere un token API Kaggle. In questo tutorial utilizzi un secret Kubernetes per le credenziali di Kaggle.

Firmare il contratto di consenso per la licenza

Per utilizzare Gemma, devi firmare il contratto di consenso. Segui queste istruzioni:

1. Accedi alla pagina del consenso per i modelli 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.

Genera un token di accesso

Per accedere al modello tramite Kaggle, devi disporre di 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.

   Viene scaricato un file denominato kaggle.json.

Carica il token di accesso in 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 bucket Cloud Storage

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

In Cloud Shell, esegui quanto segue:

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

Sostituisci CHECKPOINTS_BUCKET_NAME con il nome del bucket Cloud Storage che memorizza i checkpoint del modello.

Copia il modello nel bucket Cloud Storage

In Cloud Shell, esegui quanto segue:

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

Deployment del modello

Caricamento di un modello

Per caricare una risorsa Model che utilizza il contenitore Saxml, esegui il seguente gcloud ai models upload comando:

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Sostituisci quanto segue:

• PROJECT_ID: l'ID del tuo progetto Google Cloud
• LOCATION_ID: la regione in cui utilizzi Vertex AI. Tieni presente che le TPU sono disponibili solo nella regione us-west1.
• DEPLOYED_MODEL_NAME: un nome per il DeployedModel. Puoi utilizzare il nome visualizzato del Model anche per DeployedModel.

Creazione di un endpoint

Devi eseguire il deployment del modello in un endpoint prima di poterlo utilizzare per fornire predizioni online. Se stai eseguendo il deployment di un modello in un endpoint esistente, puoi saltare questo passaggio. Nell'esempio seguente viene utilizzato il comando gcloud ai endpoints create:

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

Sostituisci quanto segue:

• LOCATION_ID: la regione in cui utilizzi Vertex AI.
• ENDPOINT_NAME: il nome visualizzato dell'endpoint.

Lo strumento Google Cloud CLI potrebbe richiedere alcuni secondi per creare l'endpoint.

Esegui il deployment del modello nell'endpoint

Una volta che l'endpoint è pronto, esegui il deployment del modello.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

Sostituisci quanto segue:

• LOCATION_ID: la regione in cui utilizzi Vertex AI.
• ENDPOINT_NAME: il nome visualizzato dell'endpoint.
• DEPLOYED_MODEL_NAME: un nome per il DeployedModel. Puoi utilizzare il nome visualizzato del Model anche per DeployedModel.

Gemma 2B può essere implementato su una macchina ct5lp-hightpu-1t più piccola. In questo caso, devi specificare --platform_topology=1x1 al momento del caricamento del modello.

Lo strumento Google Cloud CLI potrebbe richiedere alcuni minuti per eseguire il deployment del modello nell'endpoint. Quando il modello viene implementato correttamente, questo comando stampa il seguente output:

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

Generazione di previsioni online dal modello di cui è stato eseguito il deployment

Per invocare il modello tramite l'endpoint di previsione di Vertex AI, formatta la richiesta di previsione utilizzando un oggetto JSON di richiesta di inferenza standard .

Nell'esempio seguente viene utilizzato il comando gcloud ai endpoints predict:

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

Sostituisci quanto segue:

• LOCATION_ID: la regione in cui utilizzi Vertex AI.
• ENDPOINT_NAME: il nome visualizzato dell'endpoint.
• instances.json ha il seguente formato: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Pulizia

Per evitare ulteriori addebiti di Vertex AI e Artifact Registry, elimina le risorse Google Cloud che hai creato durante questo tutorial:

1. Per annullare il deployment del modello dall'endpoint ed eliminare l'endpoint, esegui il seguente comando nella shell:

   ENDPOINT_ID=$(gcloud ai endpoints list \
      --region=LOCATION \
      --filter=display_name=ENDPOINT_NAME \
      --format="value(name)")
   
   DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
      --region=LOCATION \
      --format="value(deployedModels.id)")
   
   gcloud ai endpoints undeploy-model $ENDPOINT_ID \
     --region=LOCATION \
     --deployed-model-id=$DEPLOYED_MODEL_ID
   
   gcloud ai endpoints delete $ENDPOINT_ID \
      --region=LOCATION \
      --quiet
   

   Sostituisci LOCATION con la regione in cui hai creato il modello in una sezione precedente.

2. Per eliminare il modello, esegui il seguente comando nella shell:

   MODEL_ID=$(gcloud ai models list \
      --region=LOCATION \
      --filter=display_name=DEPLOYED_MODEL_NAME \
      --format="value(name)")
   
   gcloud ai models delete $MODEL_ID \
      --region=LOCATION \
      --quiet
   

   Sostituisci LOCATION con la regione in cui hai creato il modello in una sezione precedente.

Limitazioni

• In Vertex AI Prediction, le Cloud TPU sono supportate solo in us-west1. Per ulteriori informazioni, consulta Località.

Passaggi successivi

• Scopri come eseguire il deployment di altri modelli Saxml come Llama2 e GPT-J.