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 mit 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 TPU-Slices mit mehreren Hosts bereitstellen, die Pathways on Cloud-Binärdateien verwenden, den JetStream-Server mit dem MaxText-Framework starten und Inferenzanfragen mit mehreren Hosts 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 Eingabeaufforderungen mit niedriger Latenz konfiguriert.

Lernpfade in Cloud

Pathways ist eine Orchestrierungsebene für Beschleuniger im großen Maßstab. 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 über einen oder mehrere große TPU-Slices hinweg koordinieren. So werden ML-Berechnungen, die Hunderte oder Tausende von TPU-Chips umfassen, optimiert.

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- und Quantisierung sowie KV-Cache-Optimierungen, 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 Fragenbeantwortung. GKE bietet die Infrastruktur, die für das verteilte Training und die 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 die Pathways-Komponenten hostet.
  • Der Pathways resource manager verwaltet Beschleunigerressourcen und koordiniert die Zuweisung von Beschleunigern für Nutzerjobs.
  • Die Pathways client-Koordinaten werden mit der Pathways resource manager koordiniert, 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.

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

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

  • Verify 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, roles/resourcemanager.projectIamAdmin

    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 column 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.
      • 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.
      • Achten Sie darauf, dass Ihr Google Cloud -Projekt für Pathways auf der Zulassungsliste steht.

      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 des entsprechenden Modells. 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} \
  --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
  --role roles/storage.objectUser

gcloud projects add-iam-policy-binding ${PROJECT} \
  --member "serviceAccount:jetstream-pathways@${PROJECT}.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}.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

      Checkpoint-Conversion

      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}" --project "${PROJECT}" --location "${ZONE}"

      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. Sie eignet sich besonders gut für KI-/ML-Arbeitslasten, bei denen ein großes Modell auf mehrere Geräte auf mehreren Knoten verteilt 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 Prüfpunktpfad 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.
        • Für einen int8-Checkpoint sollte er ähnlich wie gs://OUTPUT_BUCKET_DIRECTORY/int8 aussehen.

        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://OUTPUT_BUCKET_DIRECTORY/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. Sie haben den JetStream-Modellserver mit MaxText und Pathways in GKE mit TPUs bereitgestellt.
      2. Ein Llama 3.1-405B-int8-Checkpoint 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 verbesserten Durchsatz und einer geringeren Latenz führen kann.

      • Prefill (Vorabfü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 \
  --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

      Checkpoint-Conversion

      Wenn Sie einen Meta Llama 2-70B-Prüfpunkt in einen MaxText-kompatiblen int8-Inferenzprüfpunkt umwandeln 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 Prüfpunkt wird bei der 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 Prüfpunktpfad 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.
        • Für einen int8-Checkpoint sollte er ähnlich wie gs://OUTPUT_BUCKET_DIRECTORY/int8 aussehen.

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