Questa guida mostra come gestire un modello LLM (Large Language Model) Gemma open model utilizzando le TPU (Tensor Processing Unit) su Vertex AI con Saxml. In questa guida, scarichi i modelli Gemma ottimizzati per le istruzioni con parametri 2B e 7B in Cloud Storage ed esegui il deployment su Vertex AI che esegue Saxml su TPU.
Sfondo
Gestendo Gemma utilizzando le TPU su Vertex AI con Saxml. Puoi usufruire di una soluzione di AI gestita che si occupa dell'infrastruttura di basso livello e offre un modo conveniente per pubblicare 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, rilasciati con una licenza aperta. Questi modelli di AI sono disponibili per l'esecuzione nelle tue 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 gestisce i modelli Paxml, JAX e PyTorch per l'inferenza. Ai fini di questo tutorial, vedremo come pubblicare Gemma su TPU più convenienti per Saxml. La configurazione delle GPU è simile. Saxml offre script per creare container per Vertex AI che utilizzeremo in questo tutorial.
TPU
Le TPU sono circuiti integrati specifici per le applicazioni (ASIC) sviluppati da Google e utilizzati per accelerare i framework di elaborazione dei dati come TensorFlow, PyTorch e JAX.
Questo tutorial mostra i modelli Gemma 2B e Gemma 7B. Vertex AI ospita questi modelli sui seguenti node pool TPU v5e single-host:
- 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
- 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.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the Vertex AI API.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the Vertex AI API.
-
In the Google Cloud console, 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.
-
Install the Google Cloud CLI.
-
If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.
-
To initialize the gcloud CLI, run the following command:
gcloud init
- Assicurati di avere una quota sufficiente per i chip TPU v5e per
Vertex AI. Per impostazione predefinita, questa quota è 0.
Per una topologia
1x1
, deve essere 1. Per2x2
, deve essere 4. Per eseguire entrambe le topologie, deve essere 5. - Crea un account Kaggle, se non ne hai già uno.
- Accedi alla pagina del consenso del modello su Kaggle.com.
- Accedi a Kaggle se non l'hai ancora fatto.
- Fai clic su Richiedi l'accesso.
- Nella sezione Scegli l'account per il consenso, seleziona Verifica tramite l'account Kaggle per utilizzare il tuo account Kaggle per il consenso.
- Accetta i Termini e condizioni del modello.
- Nel browser, vai alle impostazioni di Kaggle.
Nella sezione API, fai clic su Crea nuovo token.
Viene scaricato un file denominato
kaggle.json
.- In Cloud Shell, fai clic su > Carica. Altro
- Seleziona File e fai clic su Scegli file.
- Apri il file
kaggle.json
. - Fai clic su Carica.
- PROJECT_ID: l'ID del tuo Google Cloud progetto
- LOCATION_ID: la regione in cui utilizzi Vertex AI. Tieni presente che le TPU sono disponibili solo in us-west1.
-
DEPLOYED_MODEL_NAME: un nome per
DeployedModel
. Puoi utilizzare il nome visualizzato diModel
anche perDeployedModel
. - LOCATION_ID: la regione in cui utilizzi Vertex AI.
- ENDPOINT_NAME: Il nome visualizzato per l'endpoint.
- LOCATION_ID: la regione in cui utilizzi Vertex AI.
- ENDPOINT_NAME: Il nome visualizzato per l'endpoint.
-
DEPLOYED_MODEL_NAME: un nome per
DeployedModel
. Puoi utilizzare il nome visualizzato diModel
anche perDeployedModel
. - LOCATION_ID: la regione in cui utilizzi Vertex AI.
- ENDPOINT_NAME: Il nome visualizzato per l'endpoint.
- instances.json ha il seguente formato:
{"instances": [{"text_batch": "<your prompt>"},{...}]}
Per annullare il deployment del modello dall'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.
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.
- Su Vertex AI le Cloud TPU sono supportate solo in
us-west1
. Per ulteriori informazioni, consulta Località. - Scopri come eseguire il deployment di altri modelli Saxml come Llama2 e GPT-J.
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:
Ottenere l'accesso al modello
Tieni presente che Cloud Shell potrebbe non disporre di risorse sufficienti per scaricare i pesi del modello. In questo caso, puoi creare un'istanza di Vertex AI Workbench per eseguire questa attività.
Per accedere ai modelli Gemma per il deployment su Vertex AI, 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 Kaggle.
Firmare il contratto di consenso alla licenza
Per utilizzare Gemma devi firmare il contratto di consenso. Segui queste istruzioni:
Generare un token di accesso
Per accedere al modello tramite Kaggle, è necessario un token API di Kaggle.
Segui questi passaggi per generare un nuovo token se non ne hai già uno:
Carica il token di accesso su Cloud Shell
In Cloud Shell, puoi caricare il token API Kaggle nel tuo Google Cloud progetto:
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
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 tuo contenitore Saxml, esegui questo 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:
Creazione di un endpoint
Devi eseguire il deployment del modello in un endpoint prima di poterlo utilizzare per fornire inferenze online. 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:
La creazione dell'endpoint potrebbe richiedere alcuni secondi.
Esegui il deployment del modello sull'endpoint
Quando l'endpoint è pronto, esegui il deployment del modello sull'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:
Gemma 2B
può essere implementato su una macchina ct5lp-hightpu-1t più piccola, nel qual 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 sull'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 inferenze online dal modello di cui è stato eseguito il deployment
Per richiamare il modello tramite l'endpoint Vertex AI, formatta la richiesta di inferenza utilizzando un oggetto JSON di richiesta di inferenza standard .
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:
Pulizia
Per evitare ulteriori addebiti per Vertex AI e addebiti per Artifact Registry, elimina le risorse Google Cloud che hai creato durante questo tutorial: