Open-Source-Modelle mithilfe von TPUs in GKE mit Optimum TPU bereitstellen


In dieser Anleitung wird beschrieben, wie Sie Open Source LLM-Modelle (Large Language Model) mit Tensor Processing Units (TPUs) in der Google Kubernetes Engine (GKE) mit dem Optimum TPU-Bereitstellungs-Framework von Hugging Face bereitstellen. In dieser Anleitung laden Sie Open-Source-Modelle von Hugging Face herunter und stellen sie in einem Standard-GKE-Cluster mit einem Container bereit, in dem Optimum TPU ausgeführt wird.

Dieser Leitfaden ist ein guter Ausgangspunkt, wenn Sie bei der Bereitstellung und Zugänglichmachung Ihrer KI/ML-Arbeitslasten die detaillierte Kontrolle, Skalierbarkeit, Robustheit, Übertragbarkeit und Kosteneffizienz von verwaltetem Kubernetes benötigen.

Diese Anleitung richtet sich an Kunden von generativer KI im Hugging Face-System, neue oder bestehende Nutzer von GKE, ML-Entwickler, MLOps-Entwickler (DevOps) oder Plattformadministratoren, die daran interessiert sind, Kubernetes-Container-Orchestrierungsfunktionen zu nutzen, um LLMs bereitzustellen.

Zur Erinnerung: Sie haben mehrere Optionen für die LLM-Inferenz in Google Cloud, darunter Angebote wie Vertex AI, GKE und Google Compute Engine, in denen Sie Bereitstellungsbibliotheken wie JetStream, vLLM und andere Partnerangebote einbinden können. Sie können beispielsweise JetStream verwenden, um die neuesten Optimierungen aus dem Projekt abzurufen. Wenn Sie Hugging Face-Optionen bevorzugen, können Sie Optimum TPU verwenden.

Optimum TPU unterstützt die folgenden Funktionen:

  • Kontinuierliches Batching
  • Token-Streaming
  • Greedysuche und multinomiale Stichprobenerhebung mit Transformern

Ziele

  1. Bereiten Sie einen GKE-Standardcluster mit der empfohlenen TPU-Topologie anhand der Modelleigenschaften vor.
  2. Optimum TPU in GKE bereitstellen
  3. Verwenden Sie Optimum TPU, um die unterstützten Modelle über curl bereitzustellen.

Hinweise

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

    Go to project selector

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

  • Enable the required API.

    Enable the API

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

    Go to project selector

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

  • Enable the required API.

    Enable the API

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      Zu IAM
    2. Wählen Sie das Projekt aus.
    3. Klicken Sie auf Zugriff erlauben.
    4. Geben Sie im Feld Neue Hauptkonten Ihre Nutzer-ID ein. Dies ist in der Regel die E-Mail-Adresse eines Google-Kontos.

    5. Wählen Sie in der Liste Rolle auswählen eine Rolle aus.
    6. Wenn Sie weitere Rollen hinzufügen möchten, klicken Sie auf Weitere Rolle hinzufügen und fügen Sie weitere Rollen hinzu.
    7. Klicken Sie auf Speichern.

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell zum Verwalten von Ressourcen, die in Google Cloud gehostet werden. Die Software, die Sie für diese Anleitung benötigen, ist in Cloud Shell vorinstalliert, einschließlich kubectl und gcloud CLI.

So richten Sie Ihre Umgebung mit Cloud Shell ein:

  1. Starten Sie in der Google Cloud Console eine Cloud Shell-Sitzung. Klicken Sie dazu in der Google Cloud Console auf Symbol für die Cloud Shell-Aktivierung Cloud Shell aktivieren. Dadurch wird im unteren Bereich der Google Cloud Console eine Sitzung gestartet.

  2. Legen Sie die Standardumgebungsvariablen fest:

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export CLUSTER_NAME=CLUSTER_NAME
    export REGION=REGION_NAME
    export ZONE=ZONE
    

    Ersetzen Sie die folgenden Werte:

    • PROJECT_ID: Ihre Google Cloud-Projekt-ID.
    • CLUSTER_NAME: der Name Ihres GKE-Clusters.
    • REGION_NAME: Die Region, in der sich der GKE-Cluster, der Cloud Storage-Bucket und die TPU-Knoten befinden. Die Region enthält Zonen, in denen TPU v5e-Maschinentypen verfügbar sind, z. B. us-west1, us-west4, us-central1, us-east1, us-east5 oder europe-west4.
    • (Nur Standardcluster) ZONE: Die Zone, in der die TPU-Ressourcen verfügbar sind (z. B. us-west4-a). Bei Autopilot-Clustern müssen Sie nicht die Zone angeben, sondern nur die Region.
  3. Klonen Sie das Optimum-TPU-Repository:

    git clone https://github.com/huggingface/optimum-tpu.git
    

Zugriff auf das Modell erhalten

Sie können die Modelle Gemma 2B oder Llama3 8B verwenden. In dieser Anleitung liegt der Schwerpunkt auf diesen beiden Modellen. Optimum TPU unterstützt jedoch weitere Modelle.

Gemma 2B

Wenn Sie Zugriff auf die Gemma-Modelle für die Bereitstellung in GKE erhalten möchten, müssen Sie zuerst die Lizenzeinwilligungsvereinbarung unterzeichnen und dann ein Hugging-Face-Zugriffstoken generieren.

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.
  2. Bestätigen Sie die Einwilligung mit Ihrem Hugging Face-Konto.
  3. Akzeptieren Sie die Modellbedingungen.

Zugriffstoken erstellen

Generieren Sie ein neues Hugging Face-Token, falls Sie noch keines haben:

  1. Klicken Sie auf Profil > Einstellungen > Zugriffstokens.
  2. Klicken Sie auf Neues Token.
  3. Geben Sie einen Namen Ihrer Wahl und eine Rolle von mindestens Read an.
  4. Klicken Sie auf Generate Token (Token generieren).
  5. Kopieren Sie das Token in die Zwischenablage.

Llama3 8B

Sie müssen die Einwilligungsvereinbarung unterzeichnen, um Llama3 8b im Hugging Face-Repository verwenden zu können.

Zugriffstoken erstellen

Generieren Sie ein neues Hugging Face-Token, falls Sie noch keines haben:

  1. Klicken Sie auf Profil > Einstellungen > Zugriffstokens.
  2. Wählen Sie Neues Token aus.
  3. Geben Sie einen Namen Ihrer Wahl und eine Rolle von mindestens Read an.
  4. Wählen Sie Token generieren aus.
  5. Kopieren Sie das Token in die Zwischenablage.

GKE-Cluster erstellen

GKE Standard-Cluster mit einem CPU-Knoten erstellen:

gcloud container create CLUSTER_NAME \
    --project=PROJECT_ID \
    --num-nodes=1 \
    --location=ZONE

TPU-Knotenpool erstellen

Erstellen Sie einen v5e-TPU-Knotenpool mit 1 Knoten und 8 Chips:

gcloud container node-pools create tpunodepool \
    --location=ZONE \
    --num-nodes=1 \
    --machine-type=ct5lp-hightpu-8t \
    --cluster=CLUSTER_NAME

Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}

Container erstellen

Führen Sie den make-Befehl aus, um das Image zu erstellen:

cd optimum-tpu && make tpu-tgi

Übertragen Sie das Image per Push in Artifact Registry:

gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest

Kubernetes-Secret für Hugging Face-Anmeldedaten erstellen

Erstellen Sie ein Kubernetes-Secret, das das Hugging Face-Token enthält:

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN} \
  --dry-run=client -o yaml | kubectl apply -f -

Optimum TPU bereitstellen

Optimum TPU bereitstellen:

Gemma 2B

  1. Speichern Sie das folgende Manifest als optimum-tpu-gemma-2b-2x4.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: tgi-tpu
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: tgi-tpu
      template:
        metadata:
          labels:
            app: tgi-tpu
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 2x4
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
          containers:
          - name: tgi-tpu
            image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
            args:
            - --model-id=google/gemma-2b
            - --max-concurrent-requests=4
            - --max-input-length=32
            - --max-total-tokens=64
            - --max-batch-size=1
            securityContext:
                privileged: true
            env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            ports:
            - containerPort: 80
            resources:
              limits:
                google.com/tpu: 8
            livenessProbe:
              httpGet:
                path: /health
                port: 80
              initialDelaySeconds: 300
              periodSeconds: 120
    
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: service
    spec:
      selector:
        app: tgi-tpu
      ports:
        - name: http
          protocol: TCP
          port: 8080
          targetPort: 80
    

    Dieses Manifest beschreibt eine Optimum TPU-Bereitstellung mit einem internen Load-Balancer über TCP-Port 8080.

  2. Wenden Sie das Manifest an:

    kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml
    

Llama3 8B

  1. Speichern Sie das folgende Manifest als optimum-tpu-llama3-8b-2x4.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: tgi-tpu
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: tgi-tpu
      template:
        metadata:
          labels:
            app: tgi-tpu
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 2x4
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
          containers:
          - name: tgi-tpu
            image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
            args:
            - --model-id=meta-llama/Meta-Llama-3-8B
            - --max-concurrent-requests=4
            - --max-input-length=32
            - --max-total-tokens=64
            - --max-batch-size=1
            env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            ports:
            - containerPort: 80
            resources:
              limits:
                google.com/tpu: 8
            livenessProbe:
              httpGet:
                path: /health
                port: 80
              initialDelaySeconds: 300
              periodSeconds: 120
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: service
    spec:
      selector:
        app: tgi-tpu
      ports:
        - name: http
          protocol: TCP
          port: 8080
          targetPort: 80
    

    Dieses Manifest beschreibt eine Optimum TPU-Bereitstellung mit einem internen Load-Balancer über TCP-Port 8080.

  2. Wenden Sie das Manifest an:

    kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml
    

So rufen Sie die Logs des laufenden Deployments auf:

kubectl logs -f -l app=tgi-tpu

Die Ausgabe sollte in etwa so aussehen:

2024-07-09T22:39:34.365472Z  WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
2024-07-09T22:40:47.851405Z  INFO text_generation_router: router/src/main.rs:314: Warming up model
2024-07-09T22:40:54.559269Z  INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
2024-07-09T22:40:54.559291Z  INFO text_generation_router: router/src/main.rs:352: Connected
2024-07-09T22:40:54.559295Z  WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0

Das Modell muss vollständig heruntergeladen sein, bevor Sie mit dem nächsten Abschnitt fortfahren.

Modell bereitstellen

Richten Sie die Portweiterleitung zum Modell ein:

kubectl port-forward svc/service 8080:8080

Mithilfe von curl mit dem Modellserver interagieren

Bereitgestellte Modelle prüfen:

Verwenden Sie in einer neuen Terminalsitzung curl, um mit Ihrem Modell zu chatten:

curl 127.0.0.1:8080/generate     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}'     -H 'Content-Type: application/json'

Die Ausgabe sollte in etwa so aussehen:

{"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

Bereitgestellte Ressourcen löschen

Mit dem folgenden Befehl vermeiden Sie, dass Ihrem Google Cloud-Konto die in dieser Anleitung erstellten Ressourcen in Rechnung gestellt werden:

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

Nächste Schritte