vLLM-Inferenz auf TPUs der Version v6e

In dieser Anleitung erfahren Sie, wie Sie die vLLM-Inferenz auf v6e-TPUs ausführen. Außerdem wird gezeigt, wie Sie das Benchmark-Script für das Meta Llama-3.1 8B-Modell ausführen.

Informationen zu den ersten Schritten mit vLLM auf v6e-TPUs finden Sie in der Kurzanleitung für vLLM.

Wenn Sie GKE verwenden, lesen Sie auch den GKE-Leitfaden.

Hinweise

Sie müssen die Einwilligungsvereinbarung unterzeichnen, um die Llama3-Modellfamilie im HuggingFace-Repository verwenden zu können. Rufen Sie https://huggingface.co/meta-llama/Llama-3.1-8B auf, füllen Sie die Einwilligungserklärung aus und warten Sie, bis Sie genehmigt wurden.

TPU v6e mit 4 Chips vorbereiten:

  1. Melden Sie sich in Ihrem Google-Konto an. Wenn Sie noch kein Konto haben, melden Sie sich hier für ein neues Konto an.
  2. Wählen Sie in der Google Cloud Console auf der Seite „Projektauswahl“ ein Google Cloud -Projekt aus oder erstellen Sie ein solches Projekt.
  3. Aktivieren Sie die Abrechnung für Ihr Google Cloud -Projekt. Die Abrechnung ist für die gesamte Nutzung von Google Cloud erforderlich.
  4. Installieren Sie die gcloud alpha-Komponenten.
  5. Führen Sie den folgenden Befehl aus, um die neueste Version der gcloud-Komponenten zu installieren.

    gcloud components update
    
  6. Aktivieren Sie die TPU API mit dem folgenden gcloud-Befehl in Cloud Shell. Sie können sie auch über die Google Cloud Console aktivieren.

    gcloud services enable tpu.googleapis.com
    
  7. Erstellen Sie eine Dienstidentität für die TPU-VM.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
  8. Erstellen Sie ein TPU-Dienstkonto und gewähren Sie Zugriff auf Google Cloud- Dienste.

    Dienstkonten ermöglichen dem TPU-Dienst von Google Cloud den Zugriff auf andere Google Cloud--Dienste. Ein nutzerverwaltetes Dienstkonto wird empfohlen. Folgen Sie diesen Anleitungen, um Rollen zu erstellen und zu gewähren. Folgende Rollen sind erforderlich:

    • TPU-Administrator: Erforderlich, um eine TPU zu erstellen
    • Storage Admin: Erforderlich für den Zugriff auf Cloud Storage
    • Logs Writer: Erforderlich zum Schreiben von Protokollen mit der Logging API
    • Monitoring Metric Writer: Erforderlich zum Schreiben von Messwerten in Cloud Monitoring
  9. Authentifizieren Sie sich bei Google Cloud und konfigurieren Sie das Standardprojekt und die Standardzone für die Google Cloud CLI.

    gcloud auth login
    gcloud config set project PROJECT_ID
    gcloud config set compute/zone ZONE

Sichere Kapazität

Wenden Sie sich an Ihr Cloud TPU-Vertriebs- oder Account-Team, um ein TPU-Kontingent anzufordern und Fragen zur Kapazität zu stellen.

Cloud TPU-Umgebung bereitstellen

Sie können v6e-TPUs mit GKE, mit GKE und XPK oder als Warteschlangenressourcen bereitstellen.

Vorbereitung

  • Prüfen Sie, ob für Ihr Projekt genügend TPUS_PER_TPU_FAMILY-Kontingent vorhanden ist. Dieses gibt die maximale Anzahl von Chips an, auf die Sie in IhremGoogle Cloud -Projekt zugreifen können.
  • Diese Anleitung wurde mit der folgenden Konfiguration getestet:
    • Python 3.10 or later
    • Nightly-Softwareversionen:
      • tägliche JAX 0.4.32.dev20240912
      • nightly LibTPU 0.1.dev20240912+nightly
    • Stabile Softwareversionen:
      • JAX + JAX-Bibliothek von v0.4.35
  • Prüfen Sie, ob Ihr Projekt ein ausreichendes TPU-Kontingent für Folgendes hat:
    • TPU-VM-Kontingent
    • Kontingent für IP-Adressen
    • Hyperdisk Balanced-Kontingent
  • Nutzerberechtigungen für Projekte

TPU v6e bereitstellen

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
      --node-id TPU_NAME \
       --project PROJECT_ID \
       --zone ZONE \
       --accelerator-type v6e-4 \
       --runtime-version v2-alpha-tpuv6e \
       --service-account SERVICE_ACCOUNT

Beschreibung der Befehls-Flags

Variable Beschreibung
NODE_ID Die vom Nutzer zugewiesene ID der TPU, die erstellt wird, wenn die anstehende Ressourcenanfrage zugewiesen wird.
PROJECT_ID Google Cloud -Projektname. Verwenden Sie ein vorhandenes Projekt oder erstellen Sie ein neues.>
ZONE Welche Zonen unterstützt werden, erfahren Sie im Dokument TPU-Regionen und ‑Zonen.
ACCELERATOR_TYPE Informationen zu den unterstützten Beschleunigertypen finden Sie in der Dokumentation unter Beschleunigertypen.
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT Das ist die E-Mail-Adresse Ihres Dienstkontos. Sie finden sie in der Google Cloud Console unter „IAM“ -> „Dienstkonten“.

Beispiel: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com

Verwenden Sie die Befehle list oder describe, um den Status der in der Warteschlange befindlichen Ressource abzufragen.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

Eine vollständige Liste der Status von anstehenden Ressourcenanfragen finden Sie in der Dokumentation zu Angeforderte Ressourcen.

Über SSH eine Verbindung zur TPU herstellen

  gcloud compute tpus tpu-vm ssh TPU_NAME

Abhängigkeiten installieren

  1. Erstellen Sie ein Verzeichnis für Miniconda:

    mkdir -p ~/miniconda3
  2. Laden Sie das Miniconda-Installationsskript herunter:

    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
  3. Installieren Sie Miniconda:

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
  4. Entfernen Sie das Miniconda-Installationsskript:

    rm -rf ~/miniconda3/miniconda.sh
  5. Fügen Sie Miniconda der Variablen PATH hinzu:

    export PATH="$HOME/miniconda3/bin:$PATH"
  6. Laden Sie ~/.bashrc neu, um die Änderungen auf die Variable PATH anzuwenden:

    source ~/.bashrc
  7. Conda-Umgebung erstellen:

    conda create -n vllm python=3.10 -y
    conda activate vllm
  8. Klonen Sie das vLLM-Repository und wechseln Sie zum vLLM-Verzeichnis:

    git clone https://github.com/vllm-project/vllm.git && cd vllm
    
  9. Bereinigen Sie die vorhandenen Pakete „torch“ und „torch-xla“:

    pip uninstall torch torch-xla -y
    
  10. Installieren Sie andere Build-Abhängigkeiten:

    pip install -r requirements-tpu.txt
    VLLM_TARGET_DEVICE="tpu" python setup.py develop
    sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
    

Zugriff auf das Modell erhalten

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

  1. Klicken Sie auf Profil > Einstellungen > Zugriffstokens.
  2. Wählen Sie Neues Token aus.
  3. Geben Sie einen Namen Ihrer Wahl und eine Rolle mit mindestens Read Berechtigungen an.
  4. Wählen Sie Token generieren aus.
  5. Kopieren Sie das generierte Token in die Zwischenablage, legen Sie es als Umgebungsvariable fest und authentifizieren Sie sich mit der huggingface-cli:

    export TOKEN=YOUR_TOKEN
    git config --global credential.helper store
    huggingface-cli login --token $TOKEN

Benchmarking-Daten herunterladen

  1. Erstellen Sie ein /data-Verzeichnis und laden Sie das ShareGPT-Dataset von Hugging Face herunter.

    mkdir ~/data && cd ~/data
    wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
    

vLLM-Server starten

Mit dem folgenden Befehl werden die Modellgewichte aus dem Hugging Face Model Hub in das Verzeichnis /tmp der TPU-VM heruntergeladen, eine Reihe von Eingabeformen vorab kompiliert und die Modellkompilierung in ~/.cache/vllm/xla_cache geschrieben.

Weitere Informationen finden Sie in der vLLM-Dokumentation.

   cd ~/vllm
   vllm serve "meta-llama/Meta-Llama-3.1-8B" --download_dir /tmp --num-scheduler-steps 4 --swap-space 16 --disable-log-requests --tensor_parallel_size=4 --max-model-len=2048 &> serve.log &

vLLM-Benchmarks ausführen

Führen Sie das vLLM-Benchmarking-Script aus:

   python benchmarks/benchmark_serving.py \
       --backend vllm \
       --model "meta-llama/Meta-Llama-3.1-8B"  \
       --dataset-name sharegpt \
       --dataset-path ~/data/ShareGPT_V3_unfiltered_cleaned_split.json  \
       --num-prompts 1000

Bereinigen

Löschen Sie die TPU:

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async