Gestisci modelli aperti Gemma utilizzando le TPU su Vertex AI Prediction con Saxml

Questa guida mostra come pubblicare un modello linguistico di grandi dimensioni (LLM) Gemma con modelli aperti utilizzando Tensor Processing Unit (TPU) su Vertex AI Prediction con Saxml. In questa guida scaricherai i modelli Gemma ottimizzati per le istruzioni dei parametri 2B e 7B in Cloud Storage ed eseguirai il deployment su Vertex AI Prediction che esegue Saxml sulle TPU.

Contesto

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

Gemma

Gemma è un insieme di modelli di intelligenza artificiale (AI) generativa, leggeri e disponibili pubblicamente, con licenza aperta. Questi modelli di IA sono disponibili per l'esecuzione in applicazioni, hardware, dispositivi mobili o servizi ospitati. 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 fornisce modelli Paxml, JAX e PyTorch per l'inferenza. Per questo tutorial, tratteremo come gestire 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 ASIC (Application-Specific Integrated Circuit) appositamente sviluppati da Google per accelerare framework di elaborazione dati come TensorFlow, PyTorch e JAX.

Questo tutorial è utile per i modelli Gemma 2B e Gemma 7B. Vertex AI Prediction ospita questi modelli sui seguenti pool di nodi TPU v5e con 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. 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.

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

    Go to project selector

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

  4. Attiva Vertex AI API.

    Abilita l'API

    5.

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

    Go to project selector

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

  7. Attiva Vertex AI API.

    Abilita l'API

    8.

  8. Nella console Google Cloud, attiva Cloud Shell.

    Attiva Cloud Shell

    Nella parte inferiore della console Google Cloud viene avviata una sessione di Cloud Shell che mostra un prompt della riga di comando. Cloud Shell è un ambiente shell con Google Cloud CLI già installato e con valori già impostati per il progetto attuale. L'inizializzazione della sessione può richiedere alcuni secondi.

Questo tutorial presuppone che utilizzi Cloud Shell per interagire con Google Cloud. Se vuoi utilizzare una shell diversa da Cloud Shell, esegui questa configurazione aggiuntiva:

  1. Installa Google Cloud CLI.

  2. Per initialize gcloud CLI, esegui questo comando: 

    gcloud init
  3. Assicurati di avere quote sufficienti per 5 chip TPU v5e per Vertex AI Prediction.
  4. Crea un account Kaggle, se non ne hai già uno.

Ottieni l'accesso al modello

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

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

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

  1. Accedi alla pagina del consenso del modello su Kaggle.com.
  2. Accedi a Kaggle se non l'hai ancora fatto.
  3. Fai clic su Richiedi accesso.
  4. Nella sezione Scegli l'account per il consenso, seleziona Verifica tramite l'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.

Segui questa procedura per generare un nuovo token, se non ne hai già uno:

  1. Nel browser, vai alle impostazioni di Kaggle.

  2. Nella sezione API, fai clic su Create New Token (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 bucket Cloud Storage

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

In Cloud Shell, esegui questo comando:

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.

Copia il modello nel bucket Cloud Storage

In Cloud Shell, esegui questo comando:

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
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# 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
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

Deployment del modello

Caricamento di un modello

Per caricare una risorsa Model che utilizza il tuo container Saxml, esegui questo comando gcloud ai models upload:

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 stai utilizzando Vertex AI. Tieni presente che le TPU sono disponibili solo in us-west1.
  • DEPLOYED_MODEL_NAME: un nome per DeployedModel. Puoi utilizzare anche il nome visualizzato di Model per DeployedModel.

Creazione di un endpoint

Prima che il modello possa essere utilizzato per fornire previsioni online, devi eseguirne il deployment su un endpoint. Se esegui il deployment di un modello in un endpoint esistente, puoi saltare questo passaggio. L'esempio seguente utilizza 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 stai utilizzando 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

Quando l'endpoint è pronto, esegui il deployment del modello nell'endpoint.

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 stai utilizzando Vertex AI.
  • ENDPOINT_NAME: il nome visualizzato dell'endpoint.
  • DEPLOYED_MODEL_NAME: un nome per DeployedModel. Puoi utilizzare anche il nome visualizzato di Model per DeployedModel.

Il deployment di Gemma 2B può essere eseguito su una macchina ct5lp-hightpu-1t più piccola. In questo caso, devi specificare --platform_topology=1x1 durante il caricamento del modello.

Lo strumento Google Cloud CLI potrebbe richiedere alcuni minuti per eseguire il deployment del modello nell'endpoint. Una volta eseguito correttamente il deployment del modello, questo comando visualizza 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 richiamare il modello tramite l'endpoint di previsione di Vertex AI, formatta la richiesta di previsione utilizzando un oggetto JSON standard Inference Request .

L'esempio seguente utilizza 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 stai utilizzando Vertex AI.
  • ENDPOINT_NAME: il nome visualizzato dell'endpoint.
  • instances.json ha il seguente formato: {"instances": [{"text_batch": "<your prompt>"},{...}]}

esegui la pulizia

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

  1. Per annullare il deployment del modello nell'endpoint ed eliminare l'endpoint, esegui questo 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 questo 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

  • Su Vertex AI Prediction le Cloud TPU sono supportate solo in us-west1. Per ulteriori informazioni, consulta la sezione località.

Passaggi successivi

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