Inférence JetStream PyTorch sur les VM TPU v6e
Ce tutoriel explique comment utiliser JetStream pour diffuser des modèles PyTorch sur TPU v6e. 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). Dans ce tutoriel, vous exécutez le benchmark d'inférence pour le modèle Llama2-7B.
Avant de commencer
Préparez-vous à provisionner un TPU v6e avec quatre puces:
Suivez le guide Configurer l'environnement Cloud TPU pour configurer un projet Google Cloud , configurer la CLI Google Cloud, activer l'API Cloud TPU et vous assurer que vous avez accès à l'utilisation des Cloud TPU.
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 sur 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.
Créez des variables d'environnement
Dans Cloud Shell, créez les variables d'environnement suivantes:export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=us-central2-b export ACCELERATOR_TYPE=v6e-4 export RUNTIME_VERSION=v2-alpha-tpuv6e export SERVICE_ACCOUNT=your-service-account export QUEUED_RESOURCE_ID=your-queued-resource-id
Description des options de commande
Variable | Description |
PROJECT_ID
|
Google Cloud nom du projet. Utilisez un projet existant ou créez-en un. |
TPU_NAME
|
Nom du TPU. |
ZONE
|
Pour connaître les zones compatibles, consultez le document Régions et zones TPU. |
ACCELERATOR_TYPE
|
Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU. |
RUNTIME_VERSION
|
Version logicielle de Cloud TPU. |
SERVICE_ACCOUNT
|
Adresse e-mail de votre compte de service . Pour le trouver, accédez à la page Comptes de service dans la console. Google Cloud
Par exemple : |
QUEUED_RESOURCE_ID
|
ID de texte attribué par l'utilisateur à la requête de ressource mise en file d'attente. |
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 ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --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}
Exécuter le benchmark JetStream PyTorch Llama2-7B
Pour configurer JetStream-PyTorch, convertir les points de contrôle du modèle et exécuter le benchmark d'inférence, suivez les instructions du dépôt GitHub.
Une fois le benchmark d'inférence terminé, veillez à nettoyer les ressources TPU.
Effectuer un nettoyage
Supprimez le TPU:
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--force \
--async