Llama 2-Modelle mit SAX auf TPU v5e bereitstellen und bereitstellen


In dieser Anleitung wird Folgendes beschrieben:

  • Cloud TPU-VM zum Bereitstellen der Llama 2-Familie von Large Language Models (LLMs) erstellen, die in verschiedenen Größen (7 B, 13 B oder 70 B) verfügbar sind
  • Prüfpunkte für die Modelle vorbereiten und in SAX bereitstellen
  • Über einen HTTP-Endpunkt mit dem Modell interagieren

Die Bereitstellung für AGI-Experimente (SAX) ist ein experimentelles System, das Paxml-, JAX- und PyTorch-Modelle zur Inferenz bereitstellt. Den Code und die Dokumentation für SAX finden Sie im Saxml-Git-Repository. Die aktuelle stabile Version mit TPU v5e-Unterstützung ist v1.1.0.

SAX-Zellen

Eine SAX-Zelle (oder -Cluster) ist die Kerneinheit für die Bereitstellung Ihrer Modelle. Sie besteht aus zwei Hauptteilen:

  • Administratorserver: Dieser Server erfasst Ihre Modellserver, weist diesen Modellservern Modelle zu und unterstützt Clients dabei, den richtigen Modellserver für die Interaktion zu finden.
  • Modellserver: Auf diesen Servern wird Ihr Modell ausgeführt. Sie sind dafür zuständig, eingehende Anfragen zu verarbeiten und Antworten zu generieren.

Das folgende Diagramm zeigt ein Diagramm einer SAX-Zelle:

SAX-Zelle mit Admin-Server und Modellservern

Abbildung 1. SAX-Zelle mit Admin-Server und Modellserver.

Sie können mit einer SAX-Zelle mithilfe von Clients, die in Python, C++ oder Go geschrieben sind, oder direkt über einen HTTP-Server interagieren. Das folgende Diagramm zeigt, wie ein externer Client mit einer SAX-Zelle interagieren kann:

Ein externer Client und ein HTTP-Endpunkt, die mit einer SAX-Zelle interagieren

Abbildung 2. Laufzeitarchitektur eines externen Clients, der mit einer SAX-Zelle interagiert.

Lernziele

  • TPU-Ressourcen für die Bereitstellung einrichten
  • SAX-Cluster erstellen
  • Llama 2-Modell veröffentlichen
  • Mit dem Modell interagieren

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

  • Cloud TPU
  • Compute Engine
  • Cloud Storage

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen. Neuen Google Cloud-Nutzern steht möglicherweise eine kostenlose Testversion zur Verfügung.

Hinweise

Richten Sie Ihr Google Cloud-Projekt ein, aktivieren Sie die Cloud TPU API und erstellen Sie ein Dienstkonto. Folgen Sie dazu der Anleitung unter Cloud TPU-Umgebung einrichten.

TPU erstellen

Die folgenden Schritte zeigen, wie Sie eine TPU-VM erstellen, die Ihr Modell bereitstellt.

  1. Erstellen Sie Umgebungsvariablen:

    export PROJECT_ID=PROJECT_ID
    export ACCELERATOR_TYPE=ACCELERATOR_TYPE
    export ZONE=ZONE
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=SERVICE_ACCOUNT
    export TPU_NAME=TPU_NAME
    export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
    

    Beschreibungen der Umgebungsvariablen

    PROJECT_ID
    Die ID Ihres Google Cloud-Projekts.
    ACCELERATOR_TYPE
    Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Unterschiedliche Llama 2-Modellgrößen haben unterschiedliche TPU-Größenanforderungen:
    • 7B: v5litepod-4 oder größer
    • 13B: v5litepod-8 oder größer
    • 70B: v5litepod-16 oder größer
    ZONE
    Die Zone, in der Sie die Cloud TPU erstellen möchten.
    SERVICE_ACCOUNT
    Das Dienstkonto, das Sie an Ihre Cloud TPU anhängen möchten.
    TPU_NAME
    Der Name für Ihre Cloud TPU.
    QUEUED_RESOURCE_ID
    Eine Kennung für die Ressourcenanfrage in der Warteschlange.
  2. Legen Sie die Projekt-ID und die Zone in der aktiven Google Cloud CLI-Konfiguration fest:

    gcloud config set project $PROJECT_ID  && gcloud config set compute/zone $ZONE
    
  3. Erstellen Sie die TPU-VM:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --accelerator-type ${ACCELERATOR_TYPE} \
      --runtime-version ${RUNTIME_VERSION} \
      --service-account ${SERVICE_ACCOUNT}
    
  4. Prüfen Sie, ob die TPU aktiv ist:

    gcloud compute tpus queued-resources list --project $PROJECT_ID --zone $ZONE
    

Checkpoint-Conversion-Knoten einrichten

Zum Ausführen der LLama-Modelle in einem SAX-Cluster müssen Sie die ursprünglichen Llama-Prüfpunkte in ein SAX-kompatibles Format konvertieren.

Für die Konvertierung sind je nach Modellgröße erhebliche Arbeitsspeicherressourcen erforderlich:

Modell Maschinentyp
7 Mrd. 50–60 GB Arbeitsspeicher
13 Mrd. 120 GB Arbeitsspeicher
70 Mrd. 500–600 GB Arbeitsspeicher (Maschinentyp N2 oder M1)

Für die Modelle 7B und 13B können Sie die Konvertierung auf der TPU-VM ausführen. Für das 70B-Modell müssen Sie eine Compute Engine-Instanz mit etwa 1 TB Speicherplatz erstellen:

gcloud compute instances create INSTANCE_NAME --project=$PROJECT_ID --zone=$ZONE \
  --machine-type=n2-highmem-128 \
  --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
  --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
  --service-account=$SERVICE_ACCOUNT \
  --scopes=https://www.googleapis.com/auth/cloud-platform \
  --tags=http-server,https-server \
  --create-disk=auto-delete=yes,boot=yes,device-name=bk-workday-dlvm,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
  --no-shielded-secure-boot \
  --shielded-vtpm \
  --shielded-integrity-monitoring \
  --labels=goog-ec-src=vm_add-gcloud \
  --reservation-affinity=any

Unabhängig davon, ob Sie eine TPU- oder Compute Engine-Instanz als Conversion-Server verwenden, richten Sie den Server zum Konvertieren der Llama 2-Prüfpunkte ein:

  1. Legen Sie für das 7B- und das 13B-Modell die Umgebungsvariable für den Servernamen auf den Namen Ihrer TPU fest:

    export CONV_SERVER_NAME=$TPU_NAME
    

    Legen Sie für das 70B-Modell die Umgebungsvariable für den Servernamen auf den Namen Ihrer Compute Engine-Instanz fest:

    export CONV_SERVER_NAME=INSTANCE_NAME
    
  2. Stellen Sie über SSH eine Verbindung zum Conversion-Knoten her.

    Wenn der Konvertierungsknoten eine TPU ist, stellen Sie eine Verbindung zur TPU her:

    gcloud compute tpus tpu-vm ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    

    Wenn der Conversion-Knoten eine Compute Engine-Instanz ist, stellen Sie eine Verbindung zur Compute Engine-VM her:

    gcloud compute ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    
  3. Installieren Sie die erforderlichen Pakete auf dem Conversion-Knoten:

    sudo apt update
    sudo apt-get install python3-pip
    sudo apt-get install git-all
    
    pip3 install paxml==1.1.0
    pip3 install torch
    pip3 install jaxlib==0.4.14
    
  4. Laden Sie das Llama-Checkpoint-Konvertierungsskript herunter:

    gcloud storage cp gs://cloud-tpu-inference-public/sax-tokenizers/llama/convert_llama_ckpt.py .
    

Llama 2-Gewichte herunterladen

Bevor Sie das Modell umwandeln, müssen Sie die Llama 2-Gewichtungen herunterladen. Für diese Anleitung müssen Sie die ursprünglichen Llama 2-Gewichtungen verwenden (z. B. meta-llama/Llama-2-7b) und nicht die Gewichtungen, die für das Hugging Face Transformers-Format konvertiert wurden (z. B. meta-llama/Llama-2-7b-hf).

Wenn Sie bereits die Llama 2-Gewichtungen haben, fahren Sie mit Gewichtungen konvertieren fort.

Zum Herunterladen der Gewichtungen aus dem Hugging Face-Hub müssen Sie ein Nutzerzugriffstoken einrichten und Zugriff auf die Llama 2-Modelle anfordern. Folgen Sie der Anleitung auf der Hugging Face-Seite für das gewünschte Modell, z. B. meta-llama/Llama-2-7b, um Zugriff anzufordern.

  1. Erstellen Sie ein Verzeichnis für die Gewichtungen:

    sudo mkdir WEIGHTS_DIRECTORY
    
  2. Hol dir die Llama2-Gewichte vom Hugging Face-Hub:

    1. Installieren Sie die Hugging Face Hub-Befehlszeile:

      pip install -U "huggingface_hub[cli]"
      
    2. Wechseln Sie in das Verzeichnis für die Gewichtung:

      cd WEIGHTS_DIRECTORY
      
    3. Lade die Llama 2-Dateien herunter:

      python3
      from huggingface_hub import login
      login()
      from huggingface_hub import hf_hub_download, snapshot_download
      import os
      PATH=os.getcwd()
      snapshot_download(repo_id="meta-llama/LLAMA2_REPO", local_dir_use_symlinks=False, local_dir=PATH)
      

      Ersetzen Sie LLAMA2_REPO durch den Namen des Hugging Face-Repositorys, von dem Sie herunterladen möchten: Llama-2-7b, Llama-2-13b oder Llama-2-70b.

Gewichtungen umrechnen

Bearbeiten Sie das Konvertierungsskript und führen Sie dann das Konvertierungsskript aus, um die Modellgewichtungen umzuwandeln.

  1. Erstellen Sie ein Verzeichnis für die konvertierten Gewichtungen:

    sudo mkdir CONVERTED_WEIGHTS
    
  2. Klonen Sie das Saxml-GitHub-Repository in ein Verzeichnis, in dem Sie Lese-, Schreib- und Ausführungsberechtigungen haben:

      git clone https://github.com/google/saxml.git -b r1.1.0
    
  3. Wechseln Sie in das Verzeichnis saxml:

    cd saxml
    
  4. Öffnen Sie die Datei saxml/tools/convert_llama_ckpt.py.

  5. Ändern Sie in der Datei saxml/tools/convert_llama_ckpt.py Zeile 169 von:

    'scale': pytorch_vars[0]['layers.%d.attention_norm.weight' % (layer_idx)].type(torch.float16).numpy()
    

    An:

    'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()
    
  6. Führen Sie das Skript saxml/tools/init_cloud_vm.sh aus:

    saxml/tools/init_cloud_vm.sh
    
  7. Nur bei 70 Mrd.: Testmodus deaktivieren:

    1. Öffnen Sie die Datei saxml/server/pax/lm/params/lm_cloud.py.
    2. Ändern Sie in der Datei saxml/server/pax/lm/params/lm_cloud.py Zeile 344 von:

      return True
      

      An:

      return False
      
  8. Gewichtungen umrechnen:

    python3 saxml/tools/convert_llama_ckpt.py --base-model-path WEIGHTS_DIRECTORY \
      --pax-model-path CONVERTED_WEIGHTS \
      --model-size MODEL_SIZE
    

    Ersetzen Sie Folgendes:

    • WEIGHTS_DIRECTORY: Verzeichnis für die ursprünglichen Gewichtungen
    • CONVERTED_WEIGHTS: Zielpfad für die konvertierten Gewichtungen.
    • MODEL_SIZE: 7b, 13b oder 70b.

Checkpoint-Verzeichnis vorbereiten

Nachdem Sie die Prüfpunkte konvertiert haben, sollte das Prüfpunktverzeichnis die folgende Struktur haben:

checkpoint_00000000
  metadata/
      metadata
    state/
        mdl_vars.params.lm*/
        ...
        ...
        step/

Erstellen Sie eine leere Datei mit dem Namen commit_success.txt und speichern Sie eine Kopie davon in den Verzeichnissen checkpoint_00000000, metadata und state. Dadurch weiß SAX, dass dieser Prüfpunkt vollständig konvertiert und zum Laden bereit ist:

  1. Wechseln Sie in das Checkpoint-Verzeichnis:

    cd CONVERTED_WEIGHTS/checkpoint_00000000
    
  2. Erstellen Sie eine leere Datei mit dem Namen commit_success.txt:

    touch commit_success.txt
    
  3. Wechseln Sie zum Metadatenverzeichnis und erstellen Sie eine leere Datei mit dem Namen commit_success.txt:

    cd metadata && touch commit_success.txt
    
  4. Wechseln Sie zum Statusverzeichnis und erstellen Sie eine leere Datei mit dem Namen commit_success.txt:

    cd .. && cd state && touch commit_success.txt
    

Das Checkpoint-Verzeichnis sollte jetzt folgende Struktur haben:

checkpoint_00000000
    commit_success.txt
metadata/
    commit_success.txt
    metadata
    state/
        commit_success.txt
        mdl_vars.params.lm*/
        ...
        ...
        step/

Cloud Storage-Bucket erstellen

Sie müssen die konvertierten Prüfpunkte in einem Cloud Storage-Bucket speichern, damit sie beim Veröffentlichen des Modells verfügbar sind.

  1. Legen Sie eine Umgebungsvariable für den Namen Ihres Cloud Storage-Bucket fest:

    export GSBUCKET=BUCKET_NAME
    
  2. Erstellen Sie einen Bucket:

    gcloud storage buckets create gs://${GSBUCKET}
    
  3. Kopieren Sie die konvertierten Prüfpunktdateien in Ihren Bucket:

    gcloud storage cp -r CONVERTED_WEIGHTS/checkpoint_00000000  gs://$GSBUCKET/sax_models/llama2/SAX_LLAMA2_DIR/
    

    Ersetzen Sie SAX_LLAMA2_DIR durch den entsprechenden Wert:

    • 7B: saxml_llama27b
    • 13B: saxml_llama213b
    • 70 Mrd.: saxml_llama270b

SAX-Cluster erstellen

So erstellen Sie einen SAX-Cluster:

In einer typischen Bereitstellung würden Sie den Administratorserver auf einer Compute Engine-Instanz und den Modellserver auf einer TPU oder GPU ausführen. Für diese Anleitung stellen Sie den Admin-Server und den Modellserver auf derselben TPU v5e-Instanz bereit.

Administratorserver erstellen

Erstellen Sie den Docker-Container für den Administratorserver:

  1. Installieren Sie Docker auf dem Conversion-Server:

    sudo apt-get update
    sudo apt-get install docker.io
    
  2. Starten Sie den Docker-Container des Administratorservers:

    sudo docker run --name sax-admin-server \
      -it \
      -d \
      --rm \
      --network host \
      --env GSBUCKET=${GSBUCKET} us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
    

Sie können den Befehl docker run ohne die Option -d ausführen, um die Logs aufzurufen und dafür zu sorgen, dass der Administratorserver korrekt gestartet wird.

Modellserver erstellen

In den folgenden Abschnitten wird beschrieben, wie Sie einen Modellserver erstellen.

7b-Modell

Starten Sie den Docker-Container des Modellservers:

sudo docker run --privileged \
  -it \
  -d \
  --rm \
  --network host \
  --name "sax-model-server" \
  --env SAX_ROOT=gs://${GSBUCKET}/sax-root us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
  --sax_cell="/sax/test" \
  --port=10001 \
  --platform_chip=tpuv5e \
  --platform_topology='4'

13b-Modell

Die Konfiguration für LLaMA13BFP16TPUv5e fehlt in lm_cloud.py. Die folgenden Schritte zeigen, wie Sie lm_cloud.py aktualisieren und ein neues Docker-Image per Commit übergeben.

  1. Starten Sie den Modellserver:

    sudo docker run --privileged \
      -it \
      -d \
      --rm \
      --network host \
      --name "sax-model-server" \
      --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
      us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
      --sax_cell="/sax/test" \
      --port=10001 \
      --platform_chip=tpuv5e \
      --platform_topology='8'
    
  2. Stellen Sie eine SSH-Verbindung zum Docker-Container her:

    sudo docker exec -it sax-model-server bash
    
  3. Installieren Sie Vim im Docker-Image:

    $ apt update
    $ apt install vim
    
  4. Öffnen Sie die Datei saxml/server/pax/lm/params/lm_cloud.py. Suchen Sie nach LLaMA13B. Sie sollten den folgenden Code sehen:

    @servable_model_registry.register
    @quantization.for_transformer(quantize_on_the_fly=False)
    class LLaMA13B(BaseLLaMA):
      """13B model on a A100-40GB.
    
      April 12, 2023
      Latency = 5.06s with 128 decoded tokens. 38ms per output token.
      """
    
      NUM_LAYERS = 40
      VOCAB_SIZE = 32000
      DIMS_PER_HEAD = 128
      NUM_HEADS = 40
      MODEL_DIMS = 5120
      HIDDEN_DIMS = 13824
      ICI_MESH_SHAPE = [1, 1, 1]
    
      @property
      def test_mode(self) -> bool:
        return True
    
  5. Kommentieren oder löschen Sie die Zeile, die mit @quantization beginnt. Nach dieser Änderung sollte die Datei so aussehen:

    @servable_model_registry.register
    class LLaMA13B(BaseLLaMA):
      """13B model on a A100-40GB.
    
      April 12, 2023
      Latency = 5.06s with 128 decoded tokens. 38ms per output token.
      """
    
      NUM_LAYERS = 40
      VOCAB_SIZE = 32000
      DIMS_PER_HEAD = 128
      NUM_HEADS = 40
      MODEL_DIMS = 5120
      HIDDEN_DIMS = 13824
      ICI_MESH_SHAPE = [1, 1, 1]
    
      @property
      def test_mode(self) -> bool:
        return True
    
  6. Fügen Sie den folgenden Code hinzu, um die TPU-Konfiguration zu unterstützen.

    @servable_model_registry.register
    class LLaMA13BFP16TPUv5e(LLaMA13B):
    """13B model on TPU v5e-8.
    
    """
    
    BATCH_SIZE = [1]
    BUCKET_KEYS = [128]
    MAX_DECODE_STEPS = [32]
    ENABLE_GENERATE_STREAM = False
    
    ICI_MESH_SHAPE = [1, 1, 8]
    
    @property
    def test_mode(self) -> bool:
      return False
    
  7. Beenden Sie die SSH-Sitzung des Docker-Containers:

    exit
    
  8. Übernehmen Sie die Änderungen in ein neues Docker-Image:

    sudo docker commit sax-model-server sax-model-server:v1.1.0-mod
    
  9. Prüfen Sie, ob das neue Docker-Image erstellt wurde:

    sudo docker images
    

    Sie können das Docker-Image in der Artifact Registry Ihres Projekts veröffentlichen. In dieser Anleitung wird jedoch das lokale Image verwendet.

  10. Beenden Sie den Modellserver. Im weiteren Verlauf dieser Anleitung wird der aktualisierte Modellserver verwendet.

    sudo docker stop sax-model-server
    
  11. Starten Sie den Modellserver mit dem aktualisierten Docker-Image. Geben Sie dabei den aktualisierten Image-Namen sax-model-server:v1.1.0-mod an:

    sudo docker run --privileged \
      -it \
      -d \
      --rm \
      --network host \
      --name "sax-model-server" \
      --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
      sax-model-server:v1.1.0-mod \
      --sax_cell="/sax/test" \
      --port=10001 \
      --platform_chip=tpuv5e \
      --platform_topology='8'
    

70B-Modell

Stellen Sie über SSH eine Verbindung zu Ihrer TPU her und starten Sie den Modellserver:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
 --project ${PROJECT_ID} \
 --zone ${ZONE} \
 --worker=all \
 --command="
   gcloud auth configure-docker \
     us-docker.pkg.dev
   # Pull SAX model server image
   sudo docker pull us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
   # Run model server
   sudo docker run \
     --privileged  \
     -it \
     -d \
     --rm \
     --network host \
     --name "sax-model-server"  \
     --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
     us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
     --sax_cell="/sax/test" \
     --port=10001 \
     --platform_chip=tpuv5e \
     --platform_topology='16'
"

Logs prüfen

Prüfen Sie in den Modellserverlogs, ob der Modellserver ordnungsgemäß gestartet wurde:

docker logs -f sax-model-server

Wenn der Modellserver nicht gestartet wurde, finden Sie weitere Informationen im Abschnitt Fehlerbehebung.

Wiederholen Sie für das 70B-Modell diese Schritte für jede TPU-VM:

  1. Stellen Sie über SSH eine Verbindung zur TPU her:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --worker=WORKER_NUMBER
    

    WORKER_NUMBER ist ein 0-basierter Index, der angibt, zu welcher TPU-VM Sie eine Verbindung herstellen möchten.

  2. Logs prüfen:

    sudo docker logs -f sax-model-server
    

    Für drei TPU-VMs sollte angezeigt werden, dass sie mit den anderen Instanzen verbunden sind:

    I1117 00:16:07.196594 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.3:10001
    I1117 00:16:07.197484 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.87:10001
    I1117 00:16:07.199437 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.13:10001
    

    Eine der TPU-VMs sollte über Logs verfügen, aus denen hervorgeht, dass der Modellserver gestartet wird:

    I1115 04:01:29.479170 139974275995200 model_service_base.py:867] Started joining SAX cell /sax/test
    ERROR: logging before flag.Parse: I1115 04:01:31.479794       1 location.go:141] Calling Join due to address update
    ERROR: logging before flag.Parse: I1115 04:01:31.814721       1 location.go:155] Joined 10.182.0.44:10000
    

Modell veröffentlichen

In SAX ist das Befehlszeilentool saxutil enthalten, das die Interaktion mit SAX-Modellservern vereinfacht. In dieser Anleitung verwenden Sie saxutil, um das Modell zu veröffentlichen. Eine vollständige Liste der saxutil-Befehle finden Sie in der Saxml-README-Datei.

  1. Wechseln Sie zu dem Verzeichnis, in das Sie das Saxml-GitHub-Repository geklont haben:

    cd  saxml
    
  2. Stellen Sie beim 70B-Modell eine Verbindung zu Ihrem Conversion-Server her:

    gcloud compute ssh ${CONV_SERVER_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
  3. Installieren Sie Bazel:

    sudo apt-get install bazel
    
  4. Legen Sie einen Alias für die Ausführung von saxutil mit Ihrem Cloud Storage-Bucket fest:

    alias saxutil='bazel run saxml/bin:saxutil -- --sax_root=gs://${GSBUCKET}/sax-root'
    
  5. Veröffentlichen Sie das Modell mit saxutil. Dies dauert auf einem TPU v5litepod-8 etwa 10 Minuten.

    saxutil --sax_root=gs://${GSBUCKET}/sax-root publish '/sax/test/MODEL' \
        saxml.server.pax.lm.params.lm_cloud.PARAMETERS \
        gs://${GSBUCKET}/sax_models/llama2/SAX_LLAMA2_DIR/checkpoint_00000000/ \
        1
    

    Ersetzen Sie die folgenden Variablen:

    Modellgröße Werte
    7 Mrd. MODEL: llama27b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama27b
    13 Mrd. MODEL: llama213b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA13BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama213b
    70 Mrd. MODEL: llama270b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA70BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama270b

Bereitstellung testen

Prüfen Sie mit dem Befehl saxutil ls, ob die Bereitstellung erfolgreich war:

saxutil ls /sax/test/MODEL

Eine erfolgreiche Bereitstellung sollte eine Anzahl von Replikaten haben, die größer als null ist, und in etwa so aussehen:

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| MODEL    | MODEL PATH                                            | CHECKPOINT PATH                                                       | # OF REPLICAS | (SELECTED) REPLICAADDRESS |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| llama27b | saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e | gs://${MODEL_BUCKET}/sax_models/llama2/7b/pax_7B/checkpoint_00000000/ | 1             | 10.182.0.28:10001         |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+

Die Docker-Logs für den Modellserver sehen in etwa so aus:

I1114 17:31:03.586631 140003787142720 model_service_base.py:532] Successfully loaded model for key: /sax/test/llama27b

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

Fehlerbehebung

Wenn die Bereitstellung fehlschlägt, prüfen Sie die Modellserverlogs:

sudo docker logs -f sax-model-server

Für eine erfolgreiche Bereitstellung sollten Sie die folgende Ausgabe sehen:

Successfully loaded model for key: /sax/test/llama27b

Wenn aus den Logs nicht hervorgeht, dass das Modell bereitgestellt wurde, prüfen Sie die Modellkonfiguration und den Pfad zum Modellprüfpunkt.

Antworten generieren

Mit dem saxutil-Tool können Sie Antworten auf Prompts generieren.

So generieren Sie Antworten auf eine Frage:

saxutil lm.generate -extra="temperature:0.2"  /sax/test/MODEL "Q: Who is Harry Potter's mother? A:"

Die Ausgabe sollte in etwa so aussehen:

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root' lm.generate /sax/test/llama27b 'Q: Who is Harry Potter's mother? A: `
+-------------------------------+------------+
| GENERATE                      | SCORE      |
+-------------------------------+------------+
| 1. Harry Potter's mother is   | -20.214787 |
| Lily Evans. 2. Harry Potter's |            |
| mother is Petunia Evans       |            |
| (Dursley).                    |            |
+-------------------------------+------------+

Mit dem Modell über einen Client interagieren

Das SAX-Repository enthält Clients, über die Sie mit einer SAX-Zelle interagieren können. Clients sind in C++, Python und Go verfügbar. Das folgende Beispiel zeigt, wie Sie einen Python-Client erstellen.

  1. Erstellen Sie den Python-Client:

    bazel build saxml/client/python:sax.cc --compile_one_dependency
    
  2. Fügen Sie den Client zu PYTHONPATH hinzu. In diesem Beispiel wird davon ausgegangen, dass sich saxml in Ihrem Basisverzeichnis befindet:

    export PYTHONPATH=${PYTHONPATH}:$HOME/saxml/bazel-bin/saxml/client/python/
    
  3. Interagieren Sie über die Python-Shell mit SAX:

    $ python3
    Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux
    Type "help", "copyright", "credits" or "license" for more information.
    >>> import sax
    >>>
    

Über einen HTTP-Endpunkt mit dem Modell interagieren

Erstellen Sie einen HTTP-Client, um über einen HTTP-Endpunkt mit dem Modell zu interagieren:

  1. Eine Compute Engine-VM erstellen:

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    export HTTP_SERVER_NAME=HTTP_SERVER_NAME
    export SERVICE_ACCOUNT=SERVICE_ACCOUNT
    export MACHINE_TYPE=e2-standard-8
    gcloud compute instances create $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE \
      --machine-type=$MACHINE_TYPE \
      --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
      --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
      --service-account=$SERVICE_ACCOUNT \
      --scopes=https://www.googleapis.com/auth/cloud-platform \
      --tags=http-server,https-server \
      --create-disk=auto-delete=yes,boot=yes,device-name=$HTTP_SERVER_NAME,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
      --no-shielded-secure-boot \
      --shielded-vtpm \
      --shielded-integrity-monitoring \
      --labels=goog-ec-src=vm_add-gcloud \
      --reservation-affinity=any
    
  2. Stellen Sie über SSH eine Verbindung zur Compute Engine-VM her:

    gcloud compute ssh $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    
  3. Klonen Sie das GitHub-Repository für KI in GKE:

    git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git
    
  4. Wechseln Sie in das HTTP-Serververzeichnis:

    cd ai-on-gke/tools/saxml-on-gke/httpserver
    
  5. Erstellen Sie die Dockerfile:

    docker build -f Dockerfile -t sax-http .
    
  6. Führen Sie den HTTP-Server aus:

    docker run -e SAX_ROOT=gs://${GSBUCKET}/sax-root -p 8888:8888 -it sax-http
    

Testen Sie den Endpunkt über Ihren lokalen Computer oder einen anderen Server mit Zugriff auf Port 8888 mit den folgenden Befehlen:

  1. Exportieren Sie Umgebungsvariablen für die IP-Adresse und den Port Ihres Servers:

    export LB_IP=HTTP_SERVER_EXTERNAL_IP
    export PORT=8888
    
  2. Legen Sie die JSON-Nutzlast fest, die das Modell und die Abfrage enthält:

    json_payload=$(cat  << EOF
    {
      "model": "/sax/test/MODEL",
      "query": "Example query"
    }
    EOF
    )
    
  3. Senden Sie diese Anfrage:

    curl --request POST --header "Content-type: application/json" -s $LB_IP:$PORT/generate --data "$json_payload"
    

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.

Wenn Sie mit dieser Anleitung fertig sind, führen Sie die folgenden Schritte aus, um Ihre Ressourcen zu bereinigen.

  1. Löschen Sie Ihre Cloud TPU.

    $ gcloud compute tpus tpu-vm delete $TPU_NAME  --zone $ZONE
    
  2. Löschen Sie Ihre Compute Engine-Instanz, falls Sie eine erstellt haben.

    gcloud compute instances delete INSTANCE_NAME
    
  3. Löschen Sie den Cloud Storage-Bucket und seinen Inhalt.

    gcloud storage rm --recursive gs://BUCKET_NAME
    

Nächste Schritte