LLMs mithilfe von Multi-Host-TPUs in GKE mit JetStream und Pathways bereitstellen

In diesem Leitfaden erfahren Sie, wie Sie moderne Large Language Models (LLMs) wie Llama 3.1 405B in Google Kubernetes Engine (GKE) mit Tensor Processing Units (TPUs) auf mehreren Knoten bereitstellen.

In diesem Leitfaden wird gezeigt, wie Sie mit portablen Open-Source-Technologien wie Kubernetes, JetStream, Pathways on Cloud und der LeaderWorkerSet-API (LWS) KI/ML-Arbeitslasten in GKE bereitstellen und bereitstellen können. Dabei werden die detaillierte Steuerung, Skalierbarkeit, Robustheit, Portabilität und Kosteneffizienz von GKE genutzt.

Hintergrund

Large Language Models sind immer größer geworden und passen nicht mehr auf ein einzelnes Host-TPU-Slice. Für die ML-Inferenz können Sie Pathways on Cloud verwenden, um die Inferenz auf mehreren Hosts in großem Maßstab in GKE auf mehreren miteinander verbundenen TPU-Knoten auszuführen. In dieser Anleitung erfahren Sie, wie Sie einen GKE-Cluster mit den Multi-Host-TPU-Slices bereitstellen, die Binärdateien von Pathways on Cloud verwenden, den JetStream-Server mit dem MaxText-Framework starten und Multi-Host-Inferenzanfragen stellen.

Wenn Sie ein LLM mit TPUs in GKE mit JetStream, MaxText und Pathways bereitstellen, können Sie eine robuste, produktionsreife Bereitstellungslösung mit allen Vorteilen von verwaltetem Kubernetes erstellen, einschließlich Kosteneffizienz, Skalierbarkeit und höherer Verfügbarkeit. In diesem Abschnitt werden die in dieser Anleitung verwendeten Schlüsseltechnologien beschrieben.

TPUs

TPUs sind von Google speziell entwickelte anwendungsspezifische integrierte Schaltungen (Application-Specific Integrated Circuits, ASICs), die verwendet werden, um das maschinelle Lernen und die KI-Modelle zu beschleunigen, die mit Frameworks wie TensorFlow, PyTorch und JAX erstellt wurden.

Bevor Sie TPUs in GKE verwenden, sollten Sie den folgenden Lernpfad durcharbeiten:

  1. Lernen Sie mehr über die aktuelle Verfügbarkeit von TPU-Versionen unter Cloud TPU-Systemarchitektur.
  2. TPUs in GKE

In dieser Anleitung wird das Bereitstellen des Modells Llama 3.1-405B beschrieben. GKE stellt das Modell auf TPU v6e-Knoten mit mehreren Hosts bereit. Dabei werden TPU-Topologien basierend auf den Modellanforderungen für die Bereitstellung von Prompts mit niedriger Latenz konfiguriert.

Pathways on Cloud

Pathways ist eine Orchestrierungsebene für Beschleuniger, die sich für große Mengen von Daten eignet. Pathways wurde speziell entwickelt, um die Erforschung neuer Systeme und ML-Forschungsideen zu ermöglichen und gleichzeitig die Leistung aktueller Modelle auf dem neuesten Stand zu halten. Mit Pathways kann ein einzelner JAX-Clientprozess Berechnungen auf einem oder mehreren großen TPU-Slices koordinieren. So werden ML-Berechnungen optimiert, die Hunderte oder Tausende von TPU-Chips umfassen.

JetStream

JetStream ist ein von Google entwickeltes Open-Source-Framework zur Bereitstellung von Inferenzen. JetStream ermöglicht leistungsstarke, durchsatzintensive und speicheroptimierte Inferenz auf TPUs und GPUs. JetStream bietet erweiterte Leistungsoptimierungen, einschließlich Techniken zur kontinuierlichen Batch-, KV-Cache- und Quantisierung, um die LLM-Bereitstellung zu erleichtern. JetStream ermöglicht die PyTorch/XLA- und JAX-TPU-Bereitstellung, um die Leistung zu optimieren.

MaxText

MaxText ist eine leistungsstarke, skalierbare und anpassbare JAX LLM-Implementierung, die auf Open-Source-JAX-Bibliotheken auf folgende Weise basiert: Flax, Orbax undOptax. Die nur Decoder-LLM-Implementierung von MaxText ist in Python geschrieben. Es nutzt den XLA-Compiler stark aus, um eine hohe Leistung zu erzielen, ohne benutzerdefinierte Kernel erstellen zu müssen.

Weitere Informationen zu den neuesten Modellen und Parametergrößen, die von MaxText unterstützt werden, finden Sie im Projekt-Repository MaxText.

Llama 3.1 405B

Llama 3.1 405B ist ein Large Language Model von Meta, das für eine Reihe von Natural Language Processing-Aufgaben entwickelt wurde, darunter Textgenerierung, Übersetzung und Question Answering. GKE bietet die Infrastruktur, die für die Anforderungen an verteiltes Training und Bereitstellung von Modellen dieser Größenordnung erforderlich ist.

Weitere Informationen finden Sie in der Llama-Dokumentation.

Architektur

In diesem Abschnitt wird die in dieser Anleitung verwendete GKE-Architektur beschrieben. Die Architektur umfasst einen GKE-Standardcluster, der TPUs bereitstellt und JetStream- und Pathways-Komponenten zum Bereitstellen und Bereitstellen des Modells hostet.

Das folgende Diagramm zeigt die Komponenten dieser Architektur:

Architektur des GKE-Cluster mit TPU-Knotenpool mit mehreren Hosts, der die JetStream- und Pathways-Komponenten enthält.

Diese Architektur umfasst die folgenden Komponenten:

  • Einen regionalen GKE Standard.
  • Ein TPU-Slice-Knotenpool mit mehreren Hosts, der die JetStream-Bereitstellung und Pathways-Komponenten hostet.
  • Der Pathways resource manager verwaltet Beschleunigerressourcen und koordiniert die Zuweisung von Beschleunigern für Nutzerjobs.
  • Die Pathways client-Koordinaten werden mit Pathways resource manager abgestimmt, um festzulegen, wo die kompilierten Programme zur Ausführung platziert werden.
  • Die Pathways worker wird auf Beschleuniger-Maschinen ausgeführt und führt Berechnungen durch. Anschließend werden Daten über den IFRT-Proxyserver an Ihre Arbeitslast zurückgesendet.
  • Das IFRT proxy client implementiert die OSS-API Interim Framework Runtime (IFRT) und fungiert als Kommunikationsbrücke zwischen Ihrer Arbeitslast und den Pathways-Komponenten.
  • Der IFRT proxy server empfängt Anfragen vom IFRT proxy client und leitet sie an den Pathways client weiter, um die Arbeit zu verteilen.
  • Der JetStream-Pathways-Container stellt einen JAX-basierten Inferenzserver bereit, der Inferenzanfragen empfängt und seine Ausführungsprozesse an den Pathways workers delegiert.
  • Die Dienstkomponente verteilt eingehenden Traffic auf alle JetStream HTTP-Replikate.
  • JetStream HTTP ist ein HTTP-Server, der Anfragen als Wrapper für das erforderliche Format von JetStream akzeptiert und an den GRPC-Client von JetStream sendet.

Hinweis

  • Melden Sie sich in Ihrem Google Cloud -Konto an. Wenn Sie mit Google Cloudnoch 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.

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

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

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

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

    •

  • Prüfen Sie, ob Sie die folgenden Rollen für das Projekt haben: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

    Rollen prüfen

    1. Rufen Sie in der Google Cloud Console die Seite IAM auf.

      IAM aufrufen
    2. Wählen Sie das Projekt aus.

    3. Suchen Sie in der Spalte Hauptkonto nach allen Zeilen, in denen Sie oder eine Gruppe, zu der Sie gehören, angegeben sind. Fragen Sie Ihren Administrator, zu welchen Gruppen Sie gehören.

    4. Prüfen Sie in allen Zeilen, in denen Sie angegeben oder enthalten sind, die Spalte Rolle, um zu sehen, ob die Liste der Rollen die erforderlichen Rollen enthält.

    Rollen zuweisen

    1. Rufen Sie in der Google Cloud Console die Seite IAM auf.

      IAM aufrufen
    2. Wählen Sie das Projekt aus.
    3. Klicken Sie auf Zugriffsrechte erteilen.

    4. Geben Sie im Feld Neue Hauptkonten Ihre Nutzer-ID ein. Das ist in der Regel die E‑Mail-Adresse eines Google-Kontos.

    5. Klicken Sie auf Rolle auswählen und suchen Sie nach der Rolle.
    6. Klicken Sie auf Weitere Rolle hinzufügen, wenn Sie weitere Rollen zuweisen möchten.
    7. Klicken Sie auf Speichern.
  • Prüfen Sie, ob Sie ein ausreichendes Kontingent für sechzehn TPU v6e PodSlice Lite-Chips haben. In dieser Anleitung verwenden Sie On-Demand-Instanzen.

Zugriff auf das Modell erhalten

So erhalten Sie Zugriff auf den Meta Llama 3.1-405B-Checkpoint für die Bereitstellung in GKE:

  1. Unterzeichnen Sie die Lizenz-Einwilligungsvereinbarung.
  2. Rufen Sie die Meta Llama-Downloadseite auf.
  3. Lesen und akzeptieren Sie die Nutzungsbedingungen des Modells, um die URL zum Herunterladen des Modells zu erhalten.
  4. Wenn Sie den Modell-Checkpoint herunterladen möchten, suchen Sie die Modell-ID für das entsprechende Modell. Eine Liste der unterstützten Modelle und ihrer IDs finden Sie in der llama-CLI-Dokumentation. Verwenden Sie beispielsweise Llama 3.1-405B-Instruct:bf16-mp16 für das Modell Llama 3.1-405B.

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell zum Verwalten von Ressourcen, die inGoogle Cloudgehostet 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
gcloud config set billing/quota_project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=CLUSTER_NAME
export BUCKET_NAME=BUCKET_NAME
export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
export NODE_LOCATION=NODE_LOCATION
export CLUSTER_VERSION=CLUSTER_VERSION
export MACHINE_TYPE=ct6e-standard-4t
export TPU_TYPE=v6e
export TOPOLOGY=4x4
export WORKERS_PER_SLICE=4

    Ersetzen Sie die folgenden Werte:

    • PROJECT_ID: Ihre Google Cloud Projekt-ID.
    • CLUSTER_NAME: der Name Ihres GKE-Clusters.
    • BUCKET_NAME: Der Name Ihres Cloud Storage-Buckets. Sie müssen das Präfix gs:// nicht angeben.
    • CONTROL_PLANE_LOCATION: Die Compute Engine-Region, in der sich der GKE-Cluster, der Cloud Storage-Bucket und die TPU-Knoten befinden. Die Region enthält Zonen, in denen TPU v6e-Maschinentypen verfügbar sind, z. B. us-east1, us-east5, europe-west4, asia-northeast1 oder us-south1.
    • NODE_LOCATION: Die Zone, in der die TPU-Ressourcen verfügbar sind (z. B. us-east1-d).
    • CLUSTER_VERSION: Die GKE-Version, die den gewünschten Maschinentyp unterstützen muss. Die GKE-Standardversion ist für Ihre Ziel-TPU möglicherweise nicht verfügbar. Eine Liste der für TPU-Maschinentypen verfügbaren GKE-Mindestversionen finden Sie unter TPU-Verfügbarkeit in GKE.
    • MACHINE_TYPE: der Maschinentyp „v6e“.
    • TPU_TYPE: Ein Präfix, das zum Benennen von Knotenpools verwendet wird (v6e).
    • TOPOLOGY: Die TPU v6e-Topologie.
    • WORKERS_PER_SLICE: Die Anzahl der Knoten pro Knotenpool oder TPU-Slice.

Google Cloud -Ressourcen erstellen und konfigurieren

So erstellen Sie die erforderlichen Ressourcen:

GKE-Cluster erstellen

  1. Regionalen GKE-Standardcluster erstellen:

    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${CONTROL_PLANE_LOCATION} \
    --scopes=cloud-platform \
    --machine-type=n2-standard-32

    Die Erstellung eines Clusters kann einige Minuten dauern.

    Ersetzen Sie CLUSTER_VERSION durch die entsprechende Clusterversion.

  2. Erstellen Sie einen TPU v6e-Knotenpool mit einer 4x4-Topologie und jeweils vier Knoten:

    gcloud container node-pools create multihost-np \
    --project=${PROJECT_ID} \
    --location=${CONTROL_PLANE_LOCATION} \
    --node-locations=${NODE_LOCATION} \
    --cluster=${CLUSTER_NAME} \
    --machine-type=${MACHINE_TYPE} \
    --num-nodes=${WORKERS_PER_SLICE} \
    --tpu-topology=${TOPOLOGY} \
    --scopes cloud-platform \
    --placement-type=COMPACT \
    --workload-metadata=GCE_METADATA

Dienstkonto für den Zugriff auf Storage-Objekte konfigurieren

Konfigurieren Sie ein Kubernetes-Dienstkonto so, dass es als IAM-Dienstkonto fungiert.

  1. Erstellen Sie ein IAM-Dienstkonto für Ihre Anwendung:

    gcloud iam service-accounts create jetstream-pathways

  2. Fügen Sie eine IAM-Richtlinienbindung für Ihr IAM-Dienstkonto hinzu, um Cloud Storage zu verwalten. So ermöglichen Sie Ihrem IAM-Dienstkonto den Zugriff auf den Storage-Bucket, in dem Ihr Prüfpunkt gespeichert wird:

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
  --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com" \
  --role roles/storage.objectUser

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
  --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com" \
  --role roles/storage.insightsCollectorService

  3. Kennzeichnen Sie das Kubernetes-Dienstkonto mit der E-Mail-Adresse des IAM-Dienstkontos. 

    kubectl annotate serviceaccount default \
iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com

Docker für die Authentifizierung bei Artifact Registry konfigurieren

Konfigurieren Sie Docker für die Authentifizierung bei Artifact Registry, damit die auf der Zulassungsliste stehenden Pathways-Images abgerufen werden können:

gcloud auth login
gcloud auth configure-docker

Prüfpunktkonvertierung

Wenn Sie einen Meta Llama 3.1-405B-Prüfpunkt in einen MaxText-kompatiblen int8-Inferenzprüfpunkt umwandeln möchten, führen Sie die Schritte unter Prüfpunktkonvertierung mit Llama3.1-405B aus. Für Ihre Bereitstellung wird der Prüfpunkt mit dem Flag load_parameters_path verwendet.

Cloud Storage-Bucket zum Speichern temporärer Pathways-Dateien erstellen

Erstellen Sie einen Cloud Storage-Bucket zum Speichern der temporären Dateien von Pathways, z. B. des Kompilierungscache:

export PATHWAYS_BUCKET=PATHWAYS_BUCKET
gcloud storage buckets create gs://$PATHWAYS_BUCKET

JetStream-MaxText und Pathways bereitstellen

Stellen Sie den JetStream-MaxText- und Pathways-Modellserver bereit.

Verbindung zum GKE-Cluster herstellen

gcloud container clusters get-credentials "${CLUSTER_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

LWS-API bereitstellen

LWS ist eine benutzerdefinierte Ressource, die für die Bereitstellung und Verwaltung zustandsorientierter, verteilter Anwendungen entwickelt wurde, insbesondere solcher mit einer Leader-Worker-Architektur. Es eignet sich besonders gut für KI-/ML-Arbeitslasten, bei denen ein großes Modell auf mehrere Geräte auf mehreren Knoten aufgeteilt und bereitgestellt wird.

VERSION=v0.6.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

Warten Sie, bis der LeaderWorkerSet-Controller vollständig verfügbar ist:

kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m

Die Ausgabe sollte in etwa so aussehen:

deployment.apps/lws-controller-manager condition met

Prüfen Sie, ob der LeaderWorkerSet-Controller im Namespace lws-system ausgeführt wird:

kubectl get pod -n lws-system

Die Ausgabe sollte in etwa so aussehen:

NAME                          READY   STATUS    RESTARTS    AGE
lws-controller-manager-abcd   1/1     Running   0           40s
lws-controller-manager-efgh   1/1     Running   0           40s

Arbeitslastmanifest bereitstellen

  1. Speichern Sie das folgende Manifest als jetstream-pathways-llama-3-1-405b-4x4.yaml:

    apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: jetstream-pathways
  annotations:
    leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  replicas: 1
  leaderWorkerTemplate:
    leaderTemplate:
      metadata:
        labels:
          app: jetstream-pathways
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: pathways-proxy
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
          args:
          imagePullPolicy: Always
          ports:
          - containerPort: 38681
        - name: pathways-rm
          env:
          - name: HOST_ADDRESS
            value: "$(LWS_LEADER_ADDRESS)"
          - name: TPU_SKIP_MDS_QUERY
            value: "true"
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          args:
          - --server_port=38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          - --node_type=resource_manager
          - --instance_count=1
          - --instance_type=tpuv6e:4x4
          imagePullPolicy: Always
          ports:
          - containerPort: 38677
        - name: jax-tpu
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
          env:
          - name: LOG_LEVEL
            value: "INFO"
          args:
          - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
          - model_name=llama3.1-405b
          - load_parameters_path=CHECKPOINT_PATH
          - max_prefill_predict_length=1024
          - max_target_length=2048
          - async_checkpointing=false
          - steps=1
          - ici_fsdp_parallelism=1
          - ici_autoregressive_parallelism=2
          - ici_tensor_parallelism=8
          - scan_layers=false
          - weight_dtype=bfloat16
          - per_device_batch_size=6
          - enable_single_controller=true
          - quantization=int8
          - quantize_kvcache=true
          - checkpoint_is_quantized=true
          - enable_model_warmup=true
          imagePullPolicy: Always
          ports:
          - containerPort: 9000
          startupProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 1
            initialDelaySeconds: 600
            failureThreshold: 10000
          livenessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 10
          readinessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 10
        - name: jetstream-http
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
          imagePullPolicy: Always
          ports:
          - containerPort: 8000
    size: 5
    workerTemplate:
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: worker
          args:
          - --server_port=38679
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          imagePullPolicy: Always
          ports:
          - containerPort: 38679
          resources:
            limits:
              google.com/tpu: "4"
--- 
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pathways
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000

  2. Legen Sie den Wert des Felds load_parameters_path auf den Pfad Ihres Prüfpunkts fest, der beim Konvertieren des Prüfpunkts erstellt wurde.

    • Bei einem bf16-Checkpoint sollte der Pfad ähnlich wie gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items aussehen.
    • Bei einem int8-Prüfpunkt sollte es ähnlich wie gs://OUTPUT_BUCKET_DIRECTORY/int8 sein.

    Legen Sie den Wert des Felds gcs_scratch_location auf den Pathways-Bucket fest, den Sie zuvor erstellt haben.

    perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/int8|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-llama-3-1-405b-4x4.yaml

Deployment-Manifest anwenden

Wenden Sie das Manifest an, um den Server bereitzustellen:

kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml

Der Modellserver sollte gestartet werden.

Start des Modellservers prüfen

Bei einem 405B-Modell kann es etwa 10 bis 20 Minuten dauern, bis der Prüfpunkt wiederhergestellt ist. Wenn Sie das Flag enable_model_warmup aktiviert haben, kann es auch länger dauern, bis das Modell bereit ist.

kubectl logs -f jetstream-pathways-0 -c jax-tpu

Die Ausgabe sieht etwa so aus:

2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
...
...
...
INFO:     Started server process [7]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

Llama 3.1-405b bereitstellen

Um das Modell Llama 3.1-405b bereitzustellen, richten Sie die Portweiterleitung ein:

kubectl port-forward svc/jetstream-svc 8000:8000

Mit der Portweiterleitung können Sie von außerhalb des Clusters auf den Dienst zugreifen. Sie können über den ClusterIP-Dienst von GKE auf das JetStream-Pathways-Deployment zugreifen. Die ClusterIP-Dienste sind nur innerhalb des Clusters erreichbar.

Mit dem Modell interagieren

Führen Sie in einem neuen Terminal den folgenden Befehl aus:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

Die erste Anfrage kann aufgrund der Aufwärmphase des Modells einige Sekunden dauern. Die Ausgabe sollte in etwa so aussehen:

{
    "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
}

Sie haben Folgendes erfolgreich ausgeführt:

  1. Der JetStream-Modellserver wurde mit MaxText und Pathways in GKE mit TPUs bereitgestellt.
  2. Ein Llama 3.1-405B-Checkpoint (int8) wurde unter gs://BUCKET_NAME erstellt.
  3. Das Modell wurde bereitgestellt und es wurde damit interagiert.

Disaggregierte Bereitstellung

Disaggregated Serving ist eine Technik zum Bereitstellen von LLMs, bei der die Prefill- und Decodierungsphasen auf verschiedene Hosts aufgeteilt werden. Dieser Ansatz optimiert die Ressourcennutzung, was zu einem besseren Durchsatz und einer geringeren Latenz führen kann.

  • Vorab ausfüllen: Ein Forward Pass für den Eingabe-Prompt, um den Key-Value-Cache zu initialisieren.

  • Decodieren: Ein Verfahren, bei dem Ausgabetokens inkrementell generiert werden, ein Token pro Schritt und ein KV-Cache-Wert pro Iteration.

  1. Legen Sie die Standardumgebungsvariablen fest:

    export NODE_POOL_NAME=dis-v6e-8
export NODE_POOL_SIZE=2
export MACHINE_TYPE=ct6e-standard-4t
export TOPOLOGY=2x4
export WORKERS_PER_SLICE=2

  2. Erstellen Sie zwei Knotenpools, die v6e-8-Knoten verwenden:

    for i in $(seq 1 ${NODE_POOL_SIZE}); do
  gcloud container node-pools create ${NODE_POOL_NAME}-${i}-np \
  --project=${PROJECT_ID} \
  --cluster=${CLUSTER_NAME} \
  --location=${CONTROL_PLANE_LOCATION} \
  --node-locations=${NODE_LOCATION} \
  --machine-type=${MACHINE_TYPE} \
  --num-nodes=${WORKERS_PER_SLICE} \
  --tpu-topology=${TOPOLOGY} \
  --scopes=cloud-platform \
  --workload-metadata=GCE_METADATA
done

Prüfpunktkonvertierung

Wenn Sie einen Meta Llama 2-70B-Prüfpunkt in einen MaxText-kompatiblen int8-Inferenzprüfpunkt konvertieren möchten, führen Sie die Schritte unter Prüfpunktkonvertierung mit Llama2-70B aus. Wählen Sie Llama2-70B als Modell aus, wenn Sie die Nutzungsbedingungen von Meta akzeptieren. Für Ihre Bereitstellung wird der Prüfpunkt mit dem Flag load_parameters_path verwendet.

Ersetzen Sie die folgenden Parameter in Ihrer checkpoint-job.yaml-Datei:

- --meta_url=META_URL
- --model_name=llama-2
- --model_path=Llama-2-70b-chat
- --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b

Der Checkpoint wird in Ihrer Bereitstellung mit dem Flag load_parameters_path verwendet.

JetStream Pathways mit disaggregierter Bereitstellung bereitstellen

  1. Speichern Sie das folgende Manifest als jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml:

    apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: jetstream-pathways
  annotations:
    leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
spec:
  replicas: 1
  leaderWorkerTemplate:
    subGroupPolicy:
      subGroupSize: 2
    leaderTemplate:
      metadata:
        labels:
          app: jetstream-pathways
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 2x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: pathways-proxy
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
          args:
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --server_port=38681
          - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
          - --xla_jf_auto_cross_replica_sharding=false
          - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
          - --xla_tpu_enable_windowed_einsum_for_all_gather=false
          - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
          - --xla_tpu_enable_experimental_fusion_cost_model=false
          - --xla_tpu_dot_dot_fusion_duplicated=false
          - --xla_tpu_dot_dot_fusion=true
          - --xla_jf_conv_input_fusion=true
          - --xla_jf_conv_output_fusion=true
          - --xla_tpu_rwb_fusion=false
          - --xla_tpu_copy_fusion_pad_unpad_ratio=0
          - --xla_tpu_licm_size_inflation_ratio=1
          - --xla_tpu_copy_elision_analysis_allowance=150000
          - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
          - --xla_tpu_order_dot_after_layout=true
          - --xla_jf_rematerialization_percent_shared_memory_limit=100
          - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
          - --xla_tpu_enforce_prefetch_fifo_order=false
          - --xla_tpu_prefetch_interval_picker_size_override=6000000
          - --xla_tpu_async_copy_bandwidth_scaling_factor=1
          - --xla_tpu_nd_short_transfer_max_chunks=-1
          - --xla_tpu_enable_aggressive_broadcast_priority_update=true
          - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
          - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
          - --xla_tpu_enable_copy_fusion=true
          - --xla_tpu_enable_cross_program_prefetch_freeing=false
          - --xla_tpu_enable_dot_strength_reduction=true
          - --xla_tpu_layout_use_dot_grouping=false
          - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
          - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
          - --xla_tpu_vector_load_fusion_window=1024
          - --xla_tpu_vector_store_fusion_window=256
          - --xla_jf_conv_reshape_fusion=false
          - --xla_tpu_input_conv_multi_users=false
          - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
          - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
          - --xla_tpu_dot_dot_fusion_separable_convs_only=false
          - --xla_tpu_enable_multi_level_nested_loop_fusion=true
          - --xla_tpu_nested_dot_fusion=true
          - --xla_tpu_enable_multi_level_nested_dot_fusion=false
          - --xla_jf_enable_multi_output_fusion=true
          - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
          - --xla_tpu_enable_flash_attention=true
          imagePullPolicy: Always
          ports:
          - containerPort: 38681
        - name: pathways-rm
          env:       
          - name: HOST_ADDRESS
            value: "$(LWS_LEADER_ADDRESS)"
          - name: TPU_SKIP_MDS_QUERY
            value: "true"
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          args:
          - --server_port=38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          - --node_type=resource_manager
          - --instance_count=2
          - --instance_type=tpuv6e:2x4
          imagePullPolicy: Always
          ports:
          - containerPort: 38677
        - name: jax-tpu
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
          args:
          - MaxText/configs/base.yml
          - tokenizer_path=assets/tokenizer.llama2
          - load_parameters_path=CHECKPOINT_PATH
          - max_prefill_predict_length=1024
          - max_target_length=2048
          - model_name=llama2-70b
          - ici_fsdp_parallelism=1
          - ici_autoregressive_parallelism=1
          - ici_tensor_parallelism=-1
          - scan_layers=false
          - weight_dtype=bfloat16
          - per_device_batch_size=27
          - checkpoint_is_quantized=true 
          - quantization=int8
          - quantize_kvcache=true
          - compute_axis_order=0,2,1,3
          - ar_cache_axis_order=0,2,1,3
          - stack_prefill_result_cache=True
          - inference_server=ExperimentalMaxtextDisaggregatedServer_8
          - inference_benchmark_test=True
          - enable_model_warmup=True
          env:
          - name: LOG_LEVEL
            value: "INFO"
          imagePullPolicy: Always
          securityContext:
            capabilities:
              add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
          ports: 
          - containerPort: 9000
          startupProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 1
            initialDelaySeconds: 240
            failureThreshold: 10000
          livenessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 100
          readinessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 100
        - name: jetstream-http
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
          imagePullPolicy: Always
          ports:
          - containerPort: 8000
    size: 5
    workerTemplate:
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 2x4
        containers:
        - name: worker
          args:
          - --server_port=38679
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          imagePullPolicy: Always
          ports:
          - containerPort: 38679
          resources:
            limits:
              google.com/tpu: "4"
--- 
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pathways
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000

  2. Legen Sie den Wert des Felds load_parameters_path auf den Pfad Ihres Prüfpunkts fest, der beim Konvertieren des Prüfpunkts erstellt wurde.

    • Bei einem bf16-Checkpoint sollte der Pfad ähnlich wie gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items aussehen.
    • Bei einem int8-Prüfpunkt sollte es ähnlich wie gs://OUTPUT_BUCKET_DIRECTORY/int8 sein.

    Legen Sie den Wert des Felds gcs_scratch_location auf den Pathways-Bucket fest, den Sie zuvor erstellt haben.

    perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/maxtext/llama-2-70b/int8|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml

  3. Wenden Sie das Manifest an:

    kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml

    Je nach Größe des Prüfpunkts kann es einige Zeit dauern, bis der Modellserver den Prüfpunkt wiederhergestellt hat. Bei einem 70B-Modell kann es etwa 8 Minuten dauern, bis der Prüfpunkt wiederhergestellt ist, einschließlich der Aufwärmphase des Modells. Sie können die Logs weiter beobachten, um den Zeitpunkt der Einsatzbereitschaft zu ermitteln, indem Sie den Start des Modellservers überprüfen und das Modell bereitstellen, indem Sie die Portweiterleitung einrichten, damit Sie mit dem Modell interagieren können.

Sie haben Folgendes erfolgreich ausgeführt:

  1. Der JetStream-Modellserver wurde mit MaxText und Pathways in GKE mit TPUs und disaggregierter Bereitstellung bereitgestellt.
  2. Es wurde ein Llama 2-70B-Int8-Checkpoint unter gs://BUCKET_NAME erstellt.
  3. Das Modell wurde bereitgestellt und es wurde damit interagiert.

Probleme beheben

  • Wenn Sie die Empty reply from server-Meldung erhalten, hat der Container möglicherweise die Modelldaten noch nicht ganz heruntergeladen. Prüfen Sie die Logs des Pods noch einmal auf die Connected-Meldung, die angibt, dass das Modell einsatzbereit ist.
  • Wenn die Meldung Connection refused angezeigt wird, prüfen Sie, ob die Portweiterleitung aktiv ist.

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

Führen Sie die folgenden Befehle aus und folgen Sie den Eingabeaufforderungen, um zu vermeiden, dass Ihrem Google Cloud Konto die in dieser Anleitung erstellten Ressourcen in Rechnung gestellt werden:

gcloud container clusters delete ${CLUSTER_NAME} --location=${CONTROL_PLANE_LOCATION}

gcloud iam service-accounts delete jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com

gcloud storage rm --recursive gs://${BUCKET_NAME}

Nächste Schritte