Inférence vLLM sur les TPU v6e

Ce tutoriel vous explique comment exécuter l'inférence vLLM sur les TPU v6e. Il vous montre également comment exécuter le script de référence pour le modèle Meta Llama-3.1-8B.

Pour commencer à utiliser vLLM sur les TPU v6e, consultez le guide de démarrage rapide vLLM.

Si vous utilisez GKE, consultez également le tutoriel GKE.

Avant de commencer

Vous devez signer le contrat de consentement pour utiliser la famille de modèles Llama3 dans le dépôt HuggingFace. Accédez à meta-llama/Llama-3.1-8B, remplissez le contrat d'autorisation et attendez d'être approuvé.

Préparez-vous à provisionner un TPU v6e avec quatre puces:

1. Suivez le guide Configurer l'environnement Cloud TPU pour configurer un projet Google Cloud , configurer Google Cloud CLI, activer l'API Cloud TPU et vous assurer que vous avez accès à l'utilisation des Cloud TPU.

2. Authentifiez-vous avec Google Cloud et configurez le projet et la zone par défaut pour Google Cloud CLI.

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

Sécuriser la capacité

Lorsque vous êtes prêt à sécuriser la capacité de TPU, consultez la section Quotas Cloud TPU pour en savoir plus sur les quotas Cloud TPU. Si vous avez d'autres questions sur la sécurisation de la capacité, contactez votre équipe commerciale ou chargée de votre compte Cloud TPU.

Provisionner l'environnement Cloud TPU

Vous pouvez provisionner des VM TPU avec GKE, avec GKE et XPK, ou en tant que ressources mises en file d'attente.

Prérequis

• Vérifiez que votre projet dispose d'un quota TPUS_PER_TPU_FAMILY suffisant, qui spécifie le nombre maximal de puces auxquelles vous pouvez accéder dans votre projetGoogle Cloud .
• Vérifiez que votre projet dispose d'un quota TPU suffisant pour :
  • Quota de VM TPU
  • Quota d'adresses IP
  • Quota Hyperdisk équilibré
• Autorisations des utilisateurs pour les projets
  • Si vous utilisez GKE avec XPK, consultez la section Autorisations de la console Cloud sur le compte utilisateur ou de service pour connaître les autorisations requises pour exécuter XPK.

Provisionner un TPU v6e

   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

Utilisez les commandes list ou describe pour interroger l'état de votre ressource mise en file d'attente.

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

Pour obtenir la liste complète des états des requêtes de ressources en file d'attente, consultez la documentation sur les ressources en file d'attente.

Se connecter au TPU à l'aide de SSH

  gcloud compute tpus tpu-vm ssh TPU_NAME

Installer des dépendances

1. Créez un répertoire pour Miniconda:

   mkdir -p ~/miniconda3

2. Téléchargez le script d'installation de Miniconda:

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

3. Installez Miniconda:

   bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

4. Supprimez le script d'installation Miniconda:

   rm -rf ~/miniconda3/miniconda.sh

5. Ajoutez Miniconda à votre variable PATH:

   export PATH="$HOME/miniconda3/bin:$PATH"

6. Actualisez ~/.bashrc pour appliquer les modifications à la variable PATH:

   source ~/.bashrc

7. Créez un environnement Conda:

   conda create -n vllm python=3.11 -y
   conda activate vllm

8. Clonez le dépôt vLLM et accédez au répertoire vllm:

   git clone https://github.com/vllm-project/vllm.git && cd vllm
   

9. Nettoyez les packages torch et torch-xla existants:

   pip uninstall torch torch-xla -y
   

10. Installez les autres dépendances de compilation:

    pip install -r requirements/tpu.txt
    VLLM_TARGET_DEVICE="tpu" python -m pip install --editable .
    sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
    

Accéder au modèle

Générez un nouveau jeton Hugging Face si vous n'en possédez pas déjà un:

1. Accédez à Your Profile > Settings > Access Tokens (Votre profil > Paramètres > Jetons d'accès).

2. Sélectionnez Créer un jeton.

3. Spécifiez un nom de votre choix et un rôle avec au moins les autorisations Read.

4. Sélectionnez Générer un jeton.

5. Copiez le jeton généré dans votre presse-papiers, définissez-le en tant que variable d'environnement et authentifiez-vous avec huggingface-cli:

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

Lancer le serveur vLLM

La commande suivante télécharge les poids du modèle à partir du hub de modèles Hugging Face dans le répertoire /tmp de la VM TPU, précompile une plage de formes d'entrée et écrit la compilation du modèle dans ~/.cache/vllm/xla_cache.

Pour en savoir plus, consultez la documentation sur vLLM.

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

Exécuter des benchmarks vLLM

Exécutez le script de benchmarking vLLM:

export MODEL="meta-llama/Llama-3.1-8B"
pip install pandas
pip install datasets
python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model $MODEL  \
  --dataset-name random \
  --random-input-len 1820 \
  --random-output-len 128 \
  --random-prefix-len 0

Effectuer un nettoyage

Supprimez le TPU:

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