Inférence vLLM sur les TPU v6e
Ce tutoriel vous explique comment exécuter l'inférence vLLM sur des 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 à https://huggingface.co/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:
- Connectez-vous à votre compte Google. Si vous ne l'avez pas déjà fait, créez un compte.
- Dans la console Google Cloud, sélectionnez ou créez un Google Cloud projet à partir de la page de sélection des projets.
- Activez la facturation pour votre Google Cloud projet. La facturation est obligatoire pour toute Google Cloud utilisation.
- Installez les composants gcloud alpha.
Exécutez la commande suivante pour installer la dernière version des composants
gcloud
.gcloud components update
Activez l'API TPU à l'aide de la commande
gcloud
suivante dans Cloud Shell. Vous pouvez également l'activer à partir de la console Google Cloud.gcloud services enable tpu.googleapis.com
Créez une identité de service pour la VM TPU.
gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
Créez un compte de service TPU et accordez-lui l'accès aux services Google Cloud .
Les comptes de service permettent au service Google Cloud TPU d'accéder à d'autres services Google Cloud. Un compte de service géré par l'utilisateur est recommandé. Suivez ces guides pour créer et accorder des rôles. Les rôles suivants sont nécessaires:
- Administrateur TPU: rôle nécessaire pour créer un TPU
- Administrateur de l'espace de stockage: rôle nécessaire pour accéder à Cloud Storage
- Écrivain de journaux: nécessaire pour écrire des journaux avec l'API Logging
- Rédacteur de métriques Monitoring: nécessaire pour écrire des métriques dans Cloud Monitoring
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é
Contactez votre équipe commerciale ou votre équipe chargée des comptes Cloud TPU pour demander un quota de TPU et poser des questions sur la capacité.
Provisionner l'environnement Cloud TPU
Vous pouvez provisionner des TPU v6e 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 . - Ce tutoriel a été testé avec la configuration suivante :
- Python
3.10 or later
- Versions logicielles nocturnes :
0.4.32.dev20240912
JAX- LibTPU
0.1.dev20240912+nightly
nightly
- Versions logicielles stables :
- JAX + Bibliothèque JAX de
v0.4.35
- JAX + Bibliothèque JAX de
- Python
- Vérifiez que votre projet dispose d'un quota TPU suffisant pour :
- Quota de VM TPU
- Quota d'adresses IP
- Quota Hyperdisk équilibré
- Autorisations de l'utilisateur sur le projet
- 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
Description des options de commande
Variable | Description |
NODE_ID | ID attribué par l'utilisateur du TPU créé lorsque la requête de ressource mise en file d'attente est allouée. |
PROJECT_ID | Google Cloud nom du projet. Utilisez un projet existant ou créez-en un.> |
ZONE | Pour connaître les zones compatibles, consultez le document Régions et zones TPU. |
ACCELERATOR_TYPE | Consultez la documentation sur les types d'accélérateurs pour connaître les types d'accélérateurs compatibles. |
RUNTIME_VERSION | v2-alpha-tpuv6e
|
SERVICE_ACCOUNT | Il s'agit de l'adresse e-mail de votre compte de service, que vous pouvez trouver dans la console Google Cloud -> IAM -> Comptes de service.
Exemple: tpu-service-account@<votre_ID_de_projet>.iam.gserviceaccount.com.com |
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
Créez un répertoire pour Miniconda:
mkdir -p ~/miniconda3
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
Installez Miniconda:
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
Supprimez le script d'installation Miniconda:
rm -rf ~/miniconda3/miniconda.sh
Ajoutez Miniconda à votre variable
PATH
:export PATH="$HOME/miniconda3/bin:$PATH"
Actualisez
~/.bashrc
pour appliquer les modifications à la variablePATH
:source ~/.bashrc
Créez un environnement Conda:
conda create -n vllm python=3.10 -y conda activate vllm
Clonez le dépôt vLLM et accédez au répertoire vLLM:
git clone https://github.com/vllm-project/vllm.git && cd vllm
Nettoyez les packages torch et torch-xla existants:
pip uninstall torch torch-xla -y
Installez les autres dépendances de compilation:
pip install -r requirements-tpu.txt VLLM_TARGET_DEVICE="tpu" python setup.py develop 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 :
- Cliquez sur Your Profile > Settings > Access Tokens (Votre profil > Paramètres > Jetons d'accès).
- Sélectionnez New Token (Nouveau jeton).
- Spécifiez un nom de votre choix et un rôle avec au moins les autorisations
Read
. - Sélectionnez Générer un jeton.
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
Télécharger les données d'analyse comparative
Créez un répertoire
/data
et téléchargez l'ensemble de données ShareGPT depuis Hugging Face.mkdir ~/data && cd ~/data wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
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/Meta-Llama-3.1-8B" --download_dir /tmp --num-scheduler-steps 4 --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 benchmark vLLM:
python benchmarks/benchmark_serving.py \
--backend vllm \
--model "meta-llama/Meta-Llama-3.1-8B" \
--dataset-name sharegpt \
--dataset-path ~/data/ShareGPT_V3_unfiltered_cleaned_split.json \
--num-prompts 1000
Effectuer un nettoyage
Supprimez le TPU:
gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \ --project PROJECT_ID \ --zone ZONE \ --force \ --async