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

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 : `tpu-service-account@PROJECT_ID.iam.gserviceaccount.com`
`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