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:

1. Suivez le guide Configurer l'environnement Cloud TPU pour vous assurer que vous disposez des droits d'accès appropriés pour utiliser Cloud TPU.

2. Créez une identité de service pour la VM TPU.

   gcloud alpha compute tpus tpu-vm service-identity create --zone=zone

3. Créez un compte de service TPU et accordez 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é. Vous pouvez créer un compte de service à partir de la console Google Cloud ou à l'aide de la commande gcloud.

   Créez un compte de service à l'aide de l'outil de ligne de commande gcloud:

   gcloud iam service-accounts create your-service-account-name \
   --description="your-sa-description" \
   --display-name="your-sa-display-name"
   export SERVICE_ACCOUNT_NAME=your-service-account-name

   Créez un compte de service à partir de la console Google Cloud:

   1. Accédez à la page "Comptes de service" dans la console Google Cloud.
   2. Cliquez sur Créer un compte de service.
   3. Saisissez le nom du compte de service.
   4. (Facultatif) Saisissez une description du compte de service.
   5. Cliquez sur Créer, puis continuez.
   6. Choisissez les rôles que vous souhaitez accorder au compte de service.
   7. Cliquez sur Continuer.
   8. (Facultatif) Spécifiez les utilisateurs ou les groupes autorisés à gérer le compte de service.
   9. Cliquez sur OK pour terminer la création du compte de service.

   Une fois votre compte de service créé, procédez comme suit pour accorder des rôles de compte de service.

   Les rôles suivants sont nécessaires:

   • Administrateur TPU: rôle nécessaire pour créer un TPU
   • Storage Admin: rôle nécessaire pour accéder à Cloud Storage
   • Rédacteur de journaux
   • Rédacteur de métriques Monitoring: nécessaire pour écrire des métriques dans Cloud Monitoring

   Votre administrateur doit vous accorder le rôle roles/resourcemanager.projectIamAdmin pour que vous puissiez attribuer des rôles IAM aux utilisateurs. Un utilisateur disposant du rôle roles/resourcemanager.projectIamAdmin Administrateur IAM de projet peut également attribuer ce rôle.

   Utilisez les commandes gcloud suivantes pour ajouter des rôles de compte de service:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/tpu.admin
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/storage.admin
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/logging.logWriter
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/monitoring.metricWriter

   Vous pouvez également attribuer des rôles à l'aide de la console Google Cloud.

   Dans la console Google Cloud, sélectionnez les rôles suivants:

   1. Sélectionnez votre compte de service, puis cliquez sur Ajouter un compte principal.
   2. Dans le champ Nouveaux comptes principaux, saisissez l'adresse e-mail de votre compte de service.
   3. Dans le menu déroulant Sélectionner un rôle, recherchez le rôle (par exemple, Storage Admin) et sélectionnez-le.
   4. Cliquez sur Enregistrer.

4. 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 page sur les quotas pour en savoir plus sur le système Cloud Quotas. Si vous avez d'autres questions sur la sécurisation de la capacité, contactez votre équipe commerciale ou votre équipe 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

• Ce tutoriel a été testé avec Python 3.10 ou version ultérieure.
• Vérifiez que votre projet dispose d'un quota TPUS_PER_TPU_FAMILY suffisant, qui spécifie le nombre maximal de chips auxquels 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 de projet utilisateur
  • 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 NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

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

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