Inférence JetStream PyTorch sur une VM Cloud TPU v5e

Suivez les étapes de la section Configurer l'environnement Cloud TPU pour créer un projet Google Cloud, activer l'API TPU, installer la CLI TPU et demander un quota TPU.

Suivez les étapes de la section Créer une ressource Cloud TPU à l'aide de l'API CreateNode pour créer une VM TPU définissant le paramètre --accelerator-type sur v5litepod-8.

Cloner le dépôt JetStream et installer les dépendances

1. Se connecter à la VM TPU à l'aide de SSH

   • Définissez ${TPU_NAME} sur le nom de votre TPU.
   • Définir ${PROJECT} sur votre projet Google Cloud
   • Définissez ${ZONE} sur la zone Google Cloud dans laquelle créer vos TPU.

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

2. Cloner le dépôt JetStream

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

   (Facultatif) Créez un environnement Python virtuel à l'aide de venv ou de conda, puis activez-le.

3. Exécuter le script d'installation

      cd jetstream-pytorch
      source install_everything.sh
   

Télécharger et convertir les pondérations

1. Téléchargez les poids officiels des lama sur GitHub.

2. Convertissez les pondérations.

   • Définissez ${IN_CKPOINT} sur l'emplacement contenant les pondérations des lama
   • Définir ${OUT_CKPOINT} sur un point de contrôle d'écriture d'emplacement

   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
   

Exécuter le moteur JetStream PyTorch en local

Pour exécuter le moteur JetStream PyTorch localement, définissez le chemin de tokenisation:

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

Exécuter le moteur JetStream PyTorch avec Llama 7B

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

Exécuter le moteur JetStream PyTorch avec Llama 13b

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

Exécuter le serveur JetStream

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

REMARQUE : Le paramètre --platform=tpu= doit spécifier le nombre d'appareils TPU (4 pour v4-8 et 8 pour v5lite-8). Exemple : --platform=tpu=8.

Après l'exécution de run_server.py, le moteur JetStream PyTorch est prêt à recevoir des appels gRPC.

Exécuter des benchmarks

Accédez au dossier deps/JetStream qui a été téléchargé lors de l'exécution de install_everything.sh.

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

Pour en savoir plus, consultez deps/JetStream/benchmarks/README.md.

Erreurs courantes

Si vous obtenez une erreur Unexpected keyword argument 'device', procédez comme suit:

• Désinstaller les dépendances jax et jaxlib
• Réinstaller avec source install_everything.sh

Si vous obtenez une erreur Out of memory, procédez comme suit:

• Utiliser une taille de lot plus petite
• Utiliser la quantification