Offene Gemma-Modelle mit TPUs in Vertex AI Prediction mit Saxml bereitstellen

Vorschau

Diese Funktion unterliegt den „Nutzungsbedingungen für Pre-GA-Angebote“ im Abschnitt „Allgemeine Nutzungsbedingungen“ der dienstspezifischen Nutzungsbedingungen. Pre-GA-Funktionen sind „wie besehen“ verfügbar und werden möglicherweise nur eingeschränkt unterstützt. Weitere Informationen finden Sie in den Beschreibungen der Startphase.

In dieser Anleitung erfahren Sie, wie Sie ein Large Language Model (LLM) von offenen Gemma-Modellen mit Tensor Processing Units (TPUs) in Vertex AI Prediction mit Saxml bereitstellen. In diesem Leitfaden laden Sie die abgestimmten Gemma-Modelle für 2B- und 7B-Parameteranweisungen in Cloud Storage herunter und stellen sie in Vertex AI Prediction bereit, das Saxml auf TPUs ausführt.

Hintergrund

Durch die Bereitstellung von Gemma mithilfe von TPUs in Vertex AI Prediction mit Saxml. Sie können eine verwaltete KI-Lösung nutzen, die sich um die Infrastruktur auf der untersten Ebene kümmert und eine kostengünstige Möglichkeit zur Bereitstellung von LLMs bietet. In diesem Abschnitt werden die in dieser Anleitung verwendeten Schlüsseltechnologien beschrieben.

Gemma

Gemma ist eine Reihe offen verfügbarer, einfacher und auf generativer KI basierender Modelle, die unter einer offenen Lizenz veröffentlicht wurden. Diese KI-Modelle können in Ihren Anwendungen, Geräten, Mobilgeräten oder gehosteten Diensten ausgeführt werden. Sie können die Gemma-Modelle zur Textgenerierung verwenden. Sie können diese Modelle jedoch auch für spezielle Aufgaben optimieren.

Weitere Informationen finden Sie in der Gemma-Dokumentation.

Saxml

Saxml ist ein experimentelles System, das Paxml-, JAX- und PyTorch-Modelle zur Inferenz bereitstellt. Für dieses Tutorial erklären wir, wie Sie Gemma auf TPUs bereitstellen, die für Saxml kosteneffizienter sind. Die Einrichtung für GPUs ist ähnlich. Saxml bietet Scripts zum Erstellen von Containern für Vertex AI Prediction, die wir in dieser Anleitung verwenden werden.

TPUs

TPUs sind von Google speziell entwickelte anwendungsspezifische integrierte Schaltkreise (ASICs), die verwendet werden, um Datenverarbeitungs-Frameworks wie TensorFlow, PyTorch und JAX zu beschleunigen.

Diese Anleitung bezieht sich auf die Modelle Gemma 2B und Gemma 7B. Vertex AI Prediction hostet diese Modelle in den folgenden TPU v5e-Knotenpools mit nur einem Host:

• Gemma 2B: Wird in einem TPU v5e-Knotenpool mit 1x1-Topologie gehostet, die einen TPU-Chip darstellt. Der Maschinentyp für die Knoten ist ct5lp-hightpu-1t
• Gemma 7B: in einem TPU v5e-Knotenpool mit 2x2- Topologie gehostet, die vier TPU-Chips darstellt. Der Maschinentyp für die Knoten ist ct5lp-hightpu-4t

Hinweise

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.

In dieser Anleitung wird davon ausgegangen, dass Sie Cloud Shell zur Interaktion mit Google Cloud verwenden. Wenn Sie anstelle von Cloud Shell eine andere Shell nutzen möchten, führen Sie die folgende zusätzliche Konfiguration aus:

1. Install the Google Cloud CLI.

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

   gcloud init

3. Achten Sie darauf, dass Ihr Kontingent für TPU v5e-Chips für Vertex AI Prediction ausreicht. Dieses Kontingent beträgt standardmäßig 0. Bei einer 1x1-Topologie muss dies 1 sein. Für 2x2 muss es 4 sein. Damit beide Topologien ausgeführt werden können, muss es 5 sein.
4. Erstellen Sie ein Kaggle-Konto, falls Sie noch keines haben.

Zugriff auf das Modell erhalten

Cloud Shell verfügt möglicherweise nicht über genügend Ressourcen, um Modellgewichtungen herunterladen. In diesem Fall können Sie ein Vertex AI Workbench-Instanz zum Ausführen dieser Aufgabe erstellen.

Damit Sie auf die Gemma-Modelle für die Bereitstellung in Vertex AI Prediction zugreifen können, müssen Sie sich in der Kaggle-Plattform anmelden, die Lizenz-Einwilligungsvereinbarung unterzeichnen und ein Kaggle-API-Token erwerben. In dieser Anleitung verwenden Sie ein Kubernetes Secret für die Kaggle-Anmeldedaten.

Lizenz-Einwilligungsvereinbarung unterzeichnen

Sie müssen die Einwilligungsvereinbarung unterzeichnen, um Gemma verwenden zu können. Gehen Sie dazu so vor:

1. Rufen Sie die Seite zur Modelleinwilligung auf Kaggle.com auf.
2. Melden Sie sich bei Kaggle an, falls Sie dies noch nicht getan haben.
3. Klicken Sie auf Zugriffsanfrage.
4. Wählen Sie im Abschnitt Konto zur Einwilligung die Option Über Kaggle-Konto verifizieren aus, um Ihr Kaggle-Konto für die Einwilligung zu verwenden.
5. Akzeptieren Sie die Nutzungsbedingungen des Modells.

Zugriffstoken erstellen

Für den Zugriff auf das Modell über Kaggle benötigen Sie ein Kaggle API-Token.

Führen Sie die folgenden Schritte aus, um ein neues Token zu generieren, falls Sie noch keines haben:

1. Rufen Sie in Ihrem Browser die Kaggle-Einstellungen auf.

2. Klicken Sie im Abschnitt API auf Neues Token erstellen.

   Eine Datei mit dem Namen kaggle.json wird heruntergeladen.

Zugriffstoken in Cloud Shell hochladen

In Cloud Shell können Sie das Kaggle API-Token in Ihr Google Cloud-Projekt hochladen:

1. Klicken Sie in Cloud Shell auf Mehr > Hochladen.
2. Wählen Sie "Datei" aus und klicken Sie auf Dateien auswählen.
3. Öffnen Sie die Datei kaggle.json.
4. Klicken Sie auf Hochladen.

Cloud Storage-Bucket erstellen

Erstellen Sie einen Cloud Storage-Bucket zum Speichern der Modellprüfpunkte.

Führen Sie in Cloud Shell folgenden Befehl aus:

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

Ersetzen Sie CHECKPOINTS_BUCKET_NAME durch den Namen des Cloud Storage-Buckets, in dem die Modellprüfpunkte gespeichert sind.

Modell in Cloud Storage-Bucket kopieren

Führen Sie in Cloud Shell folgenden Befehl aus:

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

Modell bereitstellen

Modell hochladen

Führen Sie folgenden gcloud ai models upload-Befehl aus, um eine Model-Ressource hochzuladen, die Ihren Saxml-Container verwendet:

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

Ersetzen Sie Folgendes:

• PROJECT_ID ist die ID Ihres Google Cloud-Projekts.
• LOCATION_ID: Die Region, in der Sie Vertex AI verwenden. TPUs sind nur in us-west1 verfügbar.
• DEPLOYED_MODEL_NAME: Ein Name für DeployedModel. Sie können auch den Anzeigenamen von Model für DeployedModel verwenden.

Endpunkt erstellen

Sie müssen das Modell auf einem Endpunkt bereitstellen, bevor es zum Bereitstellen von Onlinevorhersagen verwendet werden kann. Wenn Sie ein Modell auf einem vorhandenen Endpunkt bereitstellen, können Sie diesen Schritt überspringen. Im folgenden Beispiel wir der Befehl gcloud ai endpoints create verwendet:

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

Dabei gilt:

• LOCATION_ID: Die Region, in der Sie Vertex AI verwenden.
• ENDPOINT_NAME: Der Anzeigename für den Endpunkt.

Es kann einige Sekunden dauern, bis das Google Cloud CLI den Endpunkt erstellt.

Modell auf dem Endpunkt bereitstellen

Wenn der Endpunkt bereit ist, stellen Sie das Modell für den Endpunkt bereit.

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

Ersetzen Sie Folgendes:

• LOCATION_ID: Die Region, in der Sie Vertex AI verwenden.
• ENDPOINT_NAME: Der Anzeigename für den Endpunkt.
• DEPLOYED_MODEL_NAME: Ein Name für DeployedModel. Sie können auch den Anzeigenamen von Model für DeployedModel verwenden.

Gemma 2B kann auf einer kleineren ct5lp-hightpu-1t-Maschine bereitgestellt werden. Dann sollten Sie beim Hochladen des Modells --platform_topology=1x1 angeben.

Es kann einige Sekunden dauern, bis die Google Cloud CLI das Modell am Endpunkt bereitstellt. Wenn das Modell erfolgreich bereitgestellt wurde, gibt dieser Befehl die folgende Ausgabe aus:

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

Onlinevorhersagen aus dem bereitgestellten Modell abrufen

Zum Aufrufen des Modells über den Vertex AI Prediction-Endpunkt formatieren Sie die Vorhersageanfrage mithilfe eines Standard-JSON-Objekts für die Inferenzanfrage .

Im folgenden Beispiel wird der Befehl gcloud ai endpoints predict verwendet:

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

Dabei gilt:

• LOCATION_ID: Die Region, in der Sie Vertex AI verwenden.
• ENDPOINT_NAME: Der Anzeigename für den Endpunkt.
• instances.json hat folgendes Format: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Bereinigen

Löschen Sie die in dieser Anleitung erstellten Google Cloud-Ressourcen, damit keine weiteren Kosten für Vertex AI-Gebühren und Artifact Registry-Gebühren anfallen:

1. Führen Sie den folgenden Befehl in der Shell aus, um die Bereitstellung des Modells beim Endpunkt aufzuheben und den Endpunkt zu löschen:

   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
   

   Ersetzen Sie dabei LOCATION durch die Region, in der Sie das Modell in einem vorherigen Abschnitt erstellt haben.

2. Führen Sie den folgenden Befehl in der Shell aus, um das Modell zu löschen:

   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
   

   Ersetzen Sie dabei LOCATION durch die Region, in der Sie das Modell in einem vorherigen Abschnitt erstellt haben.

Beschränkungen

• Unter Vertex AI Prediction werden Cloud TPUs nur in us-west1 unterstützt. Weitere Informationen finden Sie unter Standorte.

Nächste Schritte

• Weitere Informationen zum Bereitstellen anderer Saxml-Modelle wie Llama2 und GPT-J