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:

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 projet Google Cloud à partir de la page de sélection du projet.
Activez la facturation pour votre projet Google Cloud . La facturation est obligatoire pour toute utilisation de Google Cloud .
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 l'accès aux services Google Cloud .

Les comptes de service permettent au service de TPU Google Cloud 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 chips auxquels 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 :
  - JAX 0.4.32.dev20240912 par nuit
  - LibTPU 0.1.dev20240912+nightly nightly
- Versions logicielles stables :
  - JAX + Bibliothèque JAX de v0.4.35
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.

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

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	Nom du projetGoogle Cloud 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@<your_project_ID>.iam.gserviceaccount.com.com
NUM_SLICES	Nombre de tranches à créer (nécessaire pour Multislice uniquement)
QUEUED_RESOURCE_ID	ID de texte attribué par l'utilisateur à la requête de ressource mise en file d'attente.
VALID_DURATION	Durée de validité de la requête de ressource mise en file d'attente.
NETWORK_NAME	Nom d'un réseau secondaire à utiliser.
NETWORK_FW_NAME	Nom d'un pare-feu réseau secondaire à utiliser.

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