JetStream PyTorch-Inferenz auf v5e-Cloud TPU-VM

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