Cette page a été traduite par l'API Cloud Translation.

Inférence JetStream PyTorch sur une VM Cloud TPU v5e

JetStream est un moteur optimisé pour le débit et la mémoire pour l'inférence de grands modèles de langage (LLM) sur les appareils XLA (TPU).

Avant de commencer

Suivez la procédure décrite dans 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 la procédure décrite dans Créer un Cloud TPU à l'aide de l'API CreateNode pour créer une VM TPU en définissant --accelerator-type sur v5litepod-8.

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

Se connecter à votre 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}
```
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 conda, puis activez-le.

Exécuter le script d'installation

   cd jetstream-pytorch
   source install_everything.sh

Télécharger et convertir les poids

Téléchargez les poids officiels de Llama sur GitHub.

Convertissez les poids.

Définissez ${IN_CKPOINT} sur l'emplacement contenant les poids de Lama.
Définissez ${OUT_CKPOINT} sur des points de contrôle d'écriture de position

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 PyTorch JetStream en local

Pour exécuter le moteur PyTorch JetStream en local, définissez le chemin d'accès au tokenizer:

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

Exécuter le moteur PyTorch JetStream 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 PyTorch JetStream 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, par exemple).--platform=tpu=8

Une fois run_server.py exécuté, le moteur PyTorch JetStream est prêt à recevoir des appels gRPC.

Exécuter des benchmarks

Accédez au dossier deps/JetStream téléchargé lorsque vous avez exécuté 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 à l'aide de 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

Effectuer un nettoyage

Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Nettoyer les dépôts GitHub

  # Clean up the JetStream repository
  rm -rf JetStream

  # Clean up the xla repository
  rm -rf xla

Nettoyer l'environnement virtuel Python
```
rm -rf .env
```
Supprimer vos ressources TPU

Pour en savoir plus, consultez Supprimer vos ressources TPU.