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 GKE-Cluster vom Typ Standard bereit. Dazu verwenden Sie einen Container, 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-Ökosystem, neue oder bestehende Nutzer von GKE, ML-Entwickler, MLOps-Entwickler (DevOps) oder Plattformadministratoren, die daran interessiert sind, Kubernetes-Container-Orchestrierungsfunktionen zum Bereitstellen von LLMs zu nutzen.

Zur Erinnerung: Für die LLM-Inferenz in Google Cloud stehen Ihnen mehrere Optionen zur Verfügung, darunter Vertex AI, GKE und Google Compute Engine. Dort können Sie Serving-Bibliotheken wie JetStream, vLLM und andere Partnerangebote einbinden. Mit JetStream können Sie beispielsweise die neuesten Optimierungen aus dem Projekt abrufen. Wenn Sie die Optionen von Hugging Face bevorzugen, können Sie Optimum TPU verwenden.

Optimum TPU unterstützt die folgenden Funktionen:

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

Zugriff auf das Modell erhalten

Sie können die Modelle Gemma 2B oder Llama3 8B verwenden. In dieser Anleitung geht es um diese beiden Modelle, aber Optimum TPU unterstützt 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

Erstellen Sie einen GKE-Standardcluster mit einem CPU-Knoten:

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

TPU-Knotenpool erstellen

So erstellen Sie einen v5e-TPU-Knotenpool mit einem Knoten und acht Chips:

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

Wenn TPU-Ressourcen verfügbar sind, stellt GKE den Knotenpool bereit. Wenn TPU-Ressourcen vorübergehend nicht verfügbar sind, wird in der Ausgabe eine GCE_STOCKOUT-Fehlermeldung angezeigt. Informationen zur Fehlerbehebung bei TPU-Engpassfehlern finden Sie unter Unzureichende TPU-Ressourcen, um die TPU-Anfrage zu erfüllen.

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

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

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

Für die Bereitstellung von Optimum TPU wird in dieser Anleitung eine Kubernetes-Bereitstellung verwendet. Ein Deployment ist ein Kubernetes-API-Objekt, mit dem Sie mehrere Replikate von Pods ausführen können, die auf die Knoten in einem Cluster verteilt sind.

Gemma 2B

  1. Speichern Sie das folgende Deployment-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=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        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=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        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"}