JetStream PyTorch-Inferenz auf v5e-Cloud TPU-VM

---

JetStream ist eine durchsatz- und speicheroptimierte Engine für LLM-Inferenzen (Large Language Models) auf XLA-Geräten (TPUs).

Hinweise

Führen Sie die Schritte unter Cloud TPU-Umgebung einrichten aus, um ein Google Cloud-Projekt zu erstellen, die TPU API zu aktivieren, die TPU-Befehlszeile zu installieren und ein TPU-Kontingent anzufordern.

Führen Sie die Schritte unter Cloud TPU mit der CreateNode API erstellen aus, um eine TPU-VM mit der Einstellung --accelerator-type auf v5litepod-8 zu erstellen.

JetStream-Repository klonen und Abhängigkeiten installieren

1. Über SSH eine Verbindung zur TPU-VM herstellen

   • Legen Sie ${TPU_NAME} auf den Namen Ihrer TPU fest.
   • ${PROJECT} auf Ihr Google Cloud-Projekt festlegen
   • ${ZONE} auf die Google Cloud-Zone festlegen, in der die TPUs erstellt werden sollen

     gcloud compute config-ssh
     gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE}
   

2. JetStream-Repository klonen

      git clone https://github.com/google/jetstream-pytorch.git
   

   (Optional) Erstellen Sie mit venv oder conda eine virtuelle Python-Umgebung und aktivieren Sie sie.

3. Installationsskript ausführen

      cd jetstream-pytorch
      source install_everything.sh
   

Gewichtungen herunterladen und konvertieren

1. Laden Sie die offiziellen Llama-Gewichte von GitHub herunter.

2. Konvertieren Sie die Gewichtungen.

   • Legen Sie ${IN_CKPOINT} auf den Ort fest, der die Lama-Gewichte enthält.
   • ${OUT_CKPOINT} auf einen Ort zum Schreiben von Checkpoints festlegen

   export input_ckpt_dir=${IN_CKPOINT}
   export output_ckpt_dir=${OUT_CKPOINT}
   export quantize=True
   python -m convert_checkpoints --input_checkpoint_dir=$input_ckpt_dir --output_checkpoint_dir=$output_ckpt_dir --quantize=$quantize
   

JetStream-PyTorch-Engine lokal ausführen

Legen Sie den Tokenizer-Pfad fest, um die JetStream PyTorch-Engine lokal auszuführen:

export tokenizer_path=${TOKENIZER_PATH} # tokenizer model file path from meta-llama

JetStream PyTorch-Engine mit Llama 7B ausführen

python run_interactive.py --size=7b --batch_size=128 --max_cache_length=2048 --quantize_weights=$quantize --quantize_kv_cache=$quantize --checkpoint_path=$output_ckpt_dir --tokenizer_path=$tokenizer_path

JetStream PyTorch-Engine mit Llama 13b ausführen

python run_interactive.py --size=13b --batch_size=64 --max_cache_length=2048 --quantize_weights=$quantize --quantize_kv_cache=$quantize --checkpoint_path=$output_ckpt_dir --tokenizer_path=$tokenizer_path

JetStream-Server ausführen

python run_server.py --param_size=7b --batch_size=128 --max_cache_length=2048 --quantize_weights=$quantize --quantize_kv_cache=$quantize --checkpoint_path=$output_ckpt_dir   --tokenizer_path=$tokenizer_path --platform=tpu=8

HINWEIS: Der Parameter --platform=tpu= muss die Anzahl der TPU-Geräte angeben (4 für v4-8 und 8 für v5lite-8). Beispiel: --platform=tpu=8.

Nach der Ausführung von run_server.py kann die JetStream PyTorch-Engine gRPC-Aufrufe empfangen.

Benchmarks ausführen

Wechseln Sie in den Ordner deps/JetStream, der beim Ausführen von install_everything.sh heruntergeladen wurde.

cd deps/JetStream
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
export dataset_path=ShareGPT_V3_unfiltered_cleaned_split.json
python benchmarks/benchmark_serving.py --tokenizer $tokenizer_path --num-prompts 2000  --dataset-path  $dataset_path --dataset sharegpt --save-request-outputs

Weitere Informationen findest du unter deps/JetStream/benchmarks/README.md.

Typische Fehler

Wenn der Fehler Unexpected keyword argument 'device' angezeigt wird, versuchen Sie Folgendes:

• Abhängigkeiten von jax und jaxlib deinstallieren
• Mit source install_everything.sh neu installieren

Wenn der Fehler Out of memory angezeigt wird, versuchen Sie Folgendes:

• Kleinere Batchgröße verwenden
• Quantisierung verwenden

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.

1. GitHub-Repositories bereinigen

     # Clean up the JetStream repository
     rm -rf JetStream
   
     # Clean up the xla repository
     rm -rf xla
   

2. Virtuelle Python-Umgebung bereinigen

   rm -rf .env
   

3. TPU-Ressourcen löschen

   Weitere Informationen finden Sie unter TPU-Ressourcen löschen.