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

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. In dieser Anleitung wird beschrieben, wie Gemma auf TPUs bereitgestellt wird, die für Saxml kostengünstiger 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 Schaltungen (Application-Specific Integrated Circuits, ASICs), die verwendet werden, um Datenverarbeitungs-Frameworks wie TensorFlow, PyTorch und JAX zu beschleunigen.

In dieser Anleitung werden die Modelle Gemma 2B und Gemma 7B bereitgestellt. 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: Wird 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. Melden Sie sich bei Ihrem Google Cloud-Konto an. Wenn Sie mit Google Cloud noch nicht vertraut sind, erstellen Sie ein Konto, um die Leistungsfähigkeit unserer Produkte in der Praxis sehen und bewerten zu können. Neukunden erhalten außerdem ein Guthaben von 300 $, um Arbeitslasten auszuführen, zu testen und bereitzustellen.

  2. Wählen Sie in der Google Cloud Console auf der Seite der Projektauswahl ein Google Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  3. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

  4. Vertex AI API and Artifact Registry API APIs aktivieren.

    Aktivieren Sie die APIs

    5.

  5. Wählen Sie in der Google Cloud Console auf der Seite der Projektauswahl ein Google Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  6. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

  7. Vertex AI API and Artifact Registry API APIs aktivieren.

    Aktivieren Sie die APIs

    8.

  8. Aktivieren Sie Cloud Shell in der Google Cloud Console.

    Cloud Shell aktivieren

    Unten in der Google Cloud Console wird eine Cloud Shell-Sitzung gestartet und eine Eingabeaufforderung angezeigt. Cloud Shell ist eine Shell-Umgebung, in der das Google Cloud CLI bereits installiert ist und Werte für Ihr aktuelles Projekt bereits festgelegt sind. Das Initialisieren der Sitzung kann einige Sekunden dauern.

In dieser Anleitung wird davon ausgegangen, dass Sie Cloud Shell für die 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. Installieren Sie die Google Cloud CLI.

  2. Führen Sie folgenden Befehl aus, um die gcloud CLI zu initialisieren: 

    gcloud init
  3. Installieren Sie Docker gemäß der Dokumentation zu Artifact Registry.
  4. Achten Sie darauf, dass Sie genügend Kontingente für 5 TPU v5e-Chips für Vertex AI Prediction haben.
  5. Erstellen Sie ein Kaggle-Konto, falls Sie noch keines haben.

Zugriff auf das Modell erhalten

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.

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 in 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 den 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
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/

Artifact Registry-Repository erstellen

Erstellen Sie ein Artifact Registry-Repository, um darin das Container-Image zu speichern, das Sie im nächsten Abschnitt erstellen.

Aktivieren Sie den Artifact Registry API-Dienst für Ihr Projekt.

gcloud services enable artifactregistry.googleapis.com

Führen Sie den folgenden Befehl in Ihrer Shell aus, um Artifact Registry-Repository zu erstellen:

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

Ersetzen Sie dabei LOCATION durch die Region, in der Artifact Registry Ihr Container-Image speichert. Später müssen Sie eine Vertex AI-Modellressource an einem regionalen Endpunkt erstellen, der zu dieser Region passt. Wählen Sie also eine Region, in der Vertex AI einen regionalen Endpunkt hat, z. B. us-west1 für TPUs.

Container-Image per Push an Artifact Registry übertragen

Der vordefinierte Saxml-Container ist unter us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest verfügbar. Kopieren Sie ihn in Ihre Artifact Registry. Konfigurieren Sie Docker für den Zugriff auf Artifact Registry. Übertragen Sie dann das Container-Image per Push in das Artifact Registry-Repository.

  1. Führen Sie den folgenden Befehl in der Shell aus, um Ihrer lokalen Docker-Installation die Berechtigung zum Übertragen an Artifact Registry per Push in der gewählten Region zu erteilen:

    gcloud auth configure-docker LOCATION-docker.pkg.dev
    • Ersetzen Sie dabei LOCATION durch die Region, in der Sie Ihr Repository erstellt haben.

  2. Führen Sie den folgenden Befehl in der Shell aus, um das Container-Image zu kopieren, das Sie gerade in Artifact Registry erstellt haben:

    docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest

  3. Führen Sie den folgenden Befehl in der Shell aus, um das Container-Image per Push zu übertragen, das Sie gerade in Artifact Registry erstellt haben:

    docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest

    Ersetzen Sie Folgendes wie im vorherigen Abschnitt:

    • LOCATION: die Region Ihres Artifact Registry-Repositorys
    • PROJECT_ID ist die ID Ihres Google Cloud-Projekts.

Modell bereitstellen

Modell hochladen

Führen Sie den 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=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-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=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-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. In diesem Fall 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 das folgende 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.

  3. Führen Sie den folgenden Befehl in der Shell aus, um Ihr Artifact Registry-Repository und das darin enthaltene Container-Image zu löschen:

    gcloud artifacts repositories delete saxml \
  --location=LOCATION \
  --quiet

    Ersetzen Sie dabei LOCATION durch die Region, in der Sie das Artifact Registry-Repository in einem vorherigen Abschnitt erstellt haben.

Beschränkungen

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

Nächste Schritte

  • Erfahren Sie, wie Sie andere Saxml-Modelle wie Llama2 und GPT-J bereitstellen.