Entraînement Cloud TPU v5e

Cloud TPU v5e est l'accélérateur d'IA de dernière génération de Google Cloud. Avec un encombrement réduit de 256 puces par pod, la version v5e est optimisée pour être le produit de plus grande valeur pour l'entraînement, l'ajustement et la diffusion des réseaux de neurones convolutifs, du texte-vers-image et des réseaux de neurones convolutifs.

Concepts, architecture système et configurations de Cloud TPU v5e

Si vous débutez avec les Cloud TPU, consultez la page d'accueil de la documentation TPU.

Les concepts généraux de Cloud TPU (par exemple, les tranches, les hôtes, les puces et les TensorCores) ainsi que l'architecture du système Cloud TPU sont décrits sur la page Architecture système de Cloud TPU.

Chaque version de Cloud TPU nécessite des types d'accélérateurs spécifiques pour l'entraînement et l'inférence. Ces types d'accélérateurs sont décrits dans la section Configurations v5e.

Inférence

L'inférence est le processus qui consiste à utiliser un modèle entraîné pour réaliser des prédictions sur de nouvelles données. Il est utilisé par le processus de diffusion.

Tranches

Une tranche représente un ensemble de puces situées toutes à l'intérieur d'un même pod connectés par des interconnexions de puces à haut débit (ICI). La version v5e présente des formes de tranche 2D. Pour connaître les formes de tranche acceptées, consultez le tableau de la section Configurations v5e.

La forme de puce et la topologie de puce font également référence aux formes de tranche.

Génération

L'inférence consiste à déployer un modèle de machine learning entraîné dans un environnement de production où il peut être utilisé pour réaliser des prédictions ou prendre des décisions. La latence et la disponibilité au niveau du service sont importantes pour la diffusion.

Hôte unique ou multi-hôte

Un hôte est un ordinateur physique (CPU) qui exécute des VM. Un hôte peut exécuter plusieurs VM à la fois.

Les segments d'application utilisant moins de huit chips utilisent au maximum un hôte. Les tranches de plus de huit puces ont accès à plusieurs hôtes et peuvent exécuter un entraînement distribué à l'aide de plusieurs hôtes. Pour en savoir plus sur les tranches et les puces, consultez la page Architecture du système TPU.

La version v5e est compatible avec l'entraînement sur plusieurs hôtes et l'inférence sur plusieurs hôtes (à l'aide de SAX).

VM TPU

Machine virtuelle exécutant Linux ayant accès aux TPU sous-jacents Pour les TPU v5e, chaque VM TPU a un accès direct à 1, 4 ou 8 puces en fonction du type d'accélérateur spécifié par l'utilisateur. Une VM TPU est également appelée nœud de calcul.

Nœud de calcul

Consultez la section VM TPU.

Démarrer

Pour en savoir plus sur le matériel TPU v5e, consultez la page Architecture du système.

Sécuriser la capacité

Contactez le service commercial Cloud pour commencer à utiliser Cloud TPU v5e pour vos charges de travail d'IA.

Préparer un projet Google Cloud

  1. Connectez-vous à votre compte Google. Si ce n'est pas déjà fait, créez un compte.
  2. Dans la console Google Cloud, sélectionnez ou créez un projet Google Cloud sur la page du sélecteur de projet.

  3. La configuration de la facturation est requise pour toute utilisation de Google Cloud. Assurez-vous donc que la facturation est activée pour votre projet.

  4. Installez les composants gcloud alpha.

  5. Si vous utilisez TPU et que vous réutilisez des composants gcloud alpha existants, mettez-les à jour pour vous assurer que les commandes et les indicateurs pertinents sont pris en charge:

    gcloud components update

  6. 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

  7. Activez le compte de service TPU.

    Les comptes de service permettent au service Cloud TPU d'accéder à d'autres services Google Cloud. L'utilisation d'un compte de service géré par l'utilisateur est une pratique recommandée dans Google Cloud. Suivez ces guides pour créer et attribuer des rôles. Les rôles suivants sont nécessaires:

    • Administrateur TPU
    • Administrateur Storage: nécessaire pour accéder à Cloud Storage.
    • Rédacteur 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.

  8. Configurez le projet et la zone.

    L'ID de votre projet est le nom de votre projet, affiché dans la console Cloud.

    export PROJECT_ID=your-project-id
export ZONE=us-west4-a

gcloud compute tpus tpu-vm service-identity create --zone=${ZONE}

gcloud auth login
gcloud config set project ${PROJECT_ID}
gcloud config set compute/zone ${ZONE}

Provisionner l'environnement Cloud TPU

Il est recommandé de provisionner des Cloud TPU v5 en tant que ressources en file d'attente à l'aide de la commande queued-resource create. Toutefois, vous pouvez également utiliser l'API Create Node (gcloud alpha compute tpus tpu-vm create) pour provisionner des Cloud TPU v5.

Créez des variables d'environnement

Définissez les variables d'environnement nécessaires pour créer des TPU.

Remplacez les variables (en rouge) de la liste suivante par les valeurs que vous utiliserez pour votre tâche d'entraînement ou d'inférence.

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d
Variable Description
PROJECT_ID Nom du projet Google Cloud
ACCELERATOR_TYPE Consultez la section Types d'accélérateurs pour connaître les types d'accélérateurs compatibles.
ZONE Toute la capacité est dans la zone us-west4-a.
RUNTIME_VERSION Utilisez la version v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT Il s'agit de l'adresse de votre compte de service que vous pouvez trouver dans la console Google Cloud -> IAM -> Comptes de service. Exemple: [email protected]
TPU_NAME ID texte attribué par l'utilisateur pour le TPU, qui est créé lorsque la requête de ressource en file d'attente est allouée.
QUEUED_RESOURCE_ID ID textuel attribué par l'utilisateur pour la requête de ressource en file d'attente. Consultez le document Ressources en file d'attente pour en savoir plus sur les ressources en file d'attente.
QUOTA_TYPE Il peut s'agir de reserved ou best-effort. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut du champ QUOTA_TYPE est on-demand. Consultez la section quotas pour plus d'informations sur les différents types de quotas acceptés par Cloud TPU.
VALID_UNTIL_DURATION Durée de validité de la requête. Consultez la section Ressources en file d'attente pour en savoir plus sur les différentes durées valides.

Créer une ressource TPU

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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

Si la ressource en file d'attente a bien été créée, l'état dans le champ response sera WAITING_FOR_RESOURCES ou FAILED. Si la ressource en file d'attente est à l'état WAITING_FOR_RESOURCES, cela signifie qu'elle a passé une validation préliminaire et qu'elle attend de la capacité. Une fois la capacité disponible, la requête passera à PROVISIONING. L'état WAITING_FOR_RESOURCES ne signifie pas que vous obtiendrez le quota alloué. Le passage de l'état WAITING_FOR_RESOURCES à l'état ACTIVE peut prendre un certain temps. Si la ressource en file d'attente est à l'état FAILED, le motif de l'échec figure dans le résultat. La requête expire si le --valid-until-duration n'est pas rempli et que l'état passe à "FAILED".

Vous pourrez accéder à votre VM TPU à l'aide de SSH une fois que votre ressource en file d'attente sera à l'état ACTIVE.

Exécutez les commandes [list](/tpu/docs/managing-tpus-tpu-vm) ou [describe](/tpu/docs/managing-tpus-tpu-vm) pour interroger l'état de la ressource en file d'attente.

gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

L'state représente l'état d'une ressource en file d'attente. Les états sont définis comme suit:

État Description
WAITING_FOR_RESOURCES La commande create de la ressource en file d'attente a été reçue et commencera le provisionnement dès que de la capacité sera disponible.
PROVISIONING Les tranches TPU sont en cours de provisionnement.
ACTIVE Tous les TPU sont provisionnés et prêts à l'emploi. Si un script de démarrage est fourni, il commence à s'exécuter sur tous les TPU lorsque l'état de la ressource en file d'attente passe à ACTIVE.
FAILED Impossible de provisionner les tranches.
SUSPENDING Un ou plusieurs secteurs sont en cours de suppression.
SUSPENDED Toutes les tranches sous-jacentes sont supprimées, mais la ressource en file d'attente reste intacte jusqu'à sa suppression explicite. Pour le moment, une ressource suspendue en file d'attente ne peut pas être réactivée et doit être supprimée.
DELETING La ressource en file d'attente est en cours de suppression.

Se connecter à la VM TPU à l'aide de SSH

La section suivante explique comment installer des binaires sur chaque VM TPU dans votre tranche TPU et comment exécuter du code. Dans ce contexte, une VM TPU est également appelée nœud de calcul.

Consultez la section Types de VM pour déterminer le nombre de VM de votre tranche.

Pour installer les binaires ou exécuter du code, connectez-vous à votre VM TPU à l'aide de la commande ssh tpu-vm.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Pour accéder à une VM TPU spécifique avec SSH, utilisez l'option --worker qui suit un index basé sur 0:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Si vos formes de secteurs comportent plus de huit puces, vous aurez plusieurs VM dans une même tranche. Dans ce cas, utilisez l'indicateur --worker=all pour exécuter l'installation sur toutes les VM TPU sans avoir à vous connecter à chacune d'elles séparément. Exemple :

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

Gérer

Toutes les commandes que vous pouvez utiliser pour gérer vos VM TPU sont décrites dans la section Gérer les TPU.

Configuration du framework

Cette section décrit le processus de configuration général pour l'entraînement de modèle personnalisé à l'aide de JAX ou PyTorch avec TPU v5e. La compatibilité avec TensorFlow est disponible dans les versions d'exécution de TPU tpu-vm-tf-2.15.0-pjrt et tpu-vm-tf-2.15.0-pod-pjrt.

Pour obtenir des instructions sur la configuration de l'inférence, consultez la page Présentation de l'inférence v5e.

Configuration pour JAX

Si vos formes de secteurs comportent plus de huit puces, vous aurez plusieurs VM dans une même tranche. Dans ce cas, vous devez utiliser l'option --worker=all pour exécuter l'installation sur toutes les VM TPU en une seule étape, sans passer par SSH pour vous connecter séparément à chacune d'elles:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les sorties présentées ici ont été générées avec une tranche v5litepod-16). Ce code vérifie que tout est correctement installé en vérifiant que JAX voit les TensorCores Cloud TPU et peut exécuter des opérations de base:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Le résultat doit ressembler à ce qui suit :

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device\_count() indique le nombre total de chips dans la tranche donnée. jax.local\_device\_count() indique le nombre de chips accessibles par une seule VM dans ce secteur.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

Le résultat doit ressembler à ce qui suit :

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Suivez les tutoriels JAX de ce document pour commencer l'entraînement dans la version 5e à l'aide de JAX.

Configuration pour PyTorch

Notez que la version v5e n'est compatible qu'avec l'environnement d'exécution PJRT et que PyTorch 2.1 et versions ultérieures utilisent PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU.

Cette section explique comment commencer à utiliser PJRT dans la version 5e avec PyTorch/XLA avec des commandes pour tous les nœuds de calcul.

Installer des dépendances

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

Si vous n'avez pas réussi à installer wheels pour torch/torch_xla/torchvision et qu'une erreur de type pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 s'affiche, passez à une version antérieure à l'aide de la commande suivante:

pip3 install setuptools==62.1.0

Exécutez un script avec PJRT:

unset LD_PRELOAD

Voici un exemple utilisant un script Python pour effectuer un calcul sur une VM v5e:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

Un résultat semblable à celui-ci doit s'afficher :

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

Suivez les tutoriels PyTorch de ce document pour commencer l'entraînement v5e à l'aide de PyTorch.

Supprimez le TPU et la ressource en file d'attente à la fin de votre session. Pour supprimer une ressource en file d'attente, supprimez la tranche, puis la ressource en file d'attente en deux étapes:

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Ces deux étapes peuvent également servir à supprimer les demandes de ressources en file d'attente qui sont à l'état FAILED.

Surveiller et profiler

Cloud TPU v5e permet la surveillance et le profilage à l'aide des mêmes méthodes que les générations précédentes de Cloud TPU. Vous pouvez consulter la page Profiler votre modèle avec les outils Cloud TPU pour en savoir plus sur le profilage, et la page Surveiller des VM Cloud TPU pour en savoir plus sur la surveillance.

Exemples JAX/FLAX

Entraîner ImageNet dans la version 5e

Ce tutoriel explique comment entraîner ImageNet dans la version 5e à l'aide de fausses données d'entrée. Si vous souhaitez utiliser des données réelles, consultez le fichier README sur GitHub.

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource QueuedResource est à l'état ACTIVE, le résultat se présente comme suit:

     state: ACTIVE

  3. Installez la dernière version de JAX et jaxlib:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Clonez le modèle ImageNet et installez la configuration requise suivante:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. Pour générer des données factices, le modèle a besoin d'informations sur les dimensions de l'ensemble de données. Vous pouvez les collecter à partir des métadonnées de l'ensemble de données ImageNet:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl http://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'

Entraîner le modèle

Une fois toutes les étapes précédentes terminées, vous pouvez entraîner le modèle.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Supprimer le TPU et la ressource en file d'attente

Supprimez le TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Modèles FLAX de Hugging Face

Les modèles de visage Hugging implémentés dans FLAX sont prêts à l'emploi sur Cloud TPU v5e. Cette section fournit des instructions pour exécuter des modèles populaires.

Entraîner ViT sur Imagenette

Ce tutoriel explique comment entraîner le modèle Vision Transformer (ViT) à partir de HuggingFace à l'aide de l'ensemble de données Fast AI imagenette sur Cloud TPU v5e.

Le modèle ViT est le premier à avoir entraîné avec succès un encodeur Transformer sur ImageNet et obtenu d'excellents résultats par rapport aux réseaux à convolution. Pour en savoir plus, consultez les ressources suivantes:

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  3. Installez JAX et sa bibliothèque:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Téléchargez le dépôt Hugging Face et les conditions requises pour l'installation:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.16.1 && pip install -r examples/flax/vision/requirements.txt'

  5. Téléchargez l'ensemble de données Imagenette:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && wget http://s3.amaz888onaws.com.fqhub.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultats de l'analyse comparative ViT

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant indique les débits avec différents types d'accélérateurs.

Type d'accélérateur V5Litepod-4 V5Litepod-16 V5Litepod-64
Époque 3 3 3
Taille du lot global 32 128 512
Débit (exemples/s) 263,40 429,34 470,71

Entraîner la diffusion sur Pokémon

Ce tutoriel explique comment entraîner le modèle de diffusion stable à partir de HuggingFace à l'aide de l'ensemble de données Pokémon sur Cloud TPU v5e.

Le modèle de diffusion stable est un modèle latent de texte à image qui génère des images photoréalistes à partir de n'importe quelle entrée de texte. Pour en savoir plus, consultez les ressources suivantes :

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  3. Installer JAX et sa bibliothèque

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Téléchargez le dépôt HuggingFace et les conditions requises pour l'installation.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.16.1 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

Supprimer le TPU et les ressources en file d'attente

Supprimez le TPU et la ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultats de l'analyse comparative de la diffusion

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant indique les débits.

Type d'accélérateur V5Litepod-4 V5Litepod-16 V5Litepod-64
Étapes d'entraînement 1500 1500 1500
Taille du lot global 32 64 128
Débit (exemples/s) 36,53 43,71 49,36

Entraîner GPT2 sur l'ensemble de données OSCAR

Ce tutoriel explique comment entraîner le modèle GPT2 à partir de HuggingFace en utilisant l'ensemble de données OSCAR sur Cloud TPU v5e.

GPT2 est un modèle Transformer pré-entraîné sur des textes bruts sans ajout manuel d'étiquettes. Il a été entraîné pour prédire le mot suivant dans des phrases. Pour en savoir plus, consultez les ressources suivantes :

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  3. Installer JAX et sa bibliothèque

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Téléchargez le dépôt HuggingFace et les conditions requises pour l'installation.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'

  5. Téléchargez les configurations pour entraîner le modèle.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultats de l'analyse comparative de GPT2

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant indique les débits.

V5Litepod-4 V5Litepod-16 V5Litepod-64
Époque 3 3 3
Taille du lot global 64 64 64
Débit (exemples/s) 74,60 72,97 72,62

PyTorch/XLA

Entraîner ResNet à l'aide de l'environnement d'exécution PJRT

PyTorch/XLA migre de XRT vers PjRt depuis PyTorch 2.0+. Voici les instructions mises à jour pour configurer la version 5e pour les charges de travail d'entraînement PyTorch/XLA.

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --{QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource QueuedResource sera à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  3. Installer des dépendances spécifiques à Torch/XLA

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'
Entraîner le modèle ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone http://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet
Résultat de l'analyse comparative

Le tableau suivant présente les débits de l'analyse comparative.

Type d'accélérateur Débit (exemples/seconde)
V5Litepod-4 4240 ex/s
V5Litepod-16 10 810 ex/s
V5Litepod-64 46 154 ex/s

Entraîner GPT2 sur v5e

Ce tutoriel explique comment exécuter GPT2 dans la version 5e du dépôt HuggingFace sur PyTorch/XLA à l'aide de l'ensemble de données wikitext.

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource QueuedResource sera à l'état ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

    state: ACTIVE

  3. Installez les dépendances Torch/xla.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get -y update
      sudo apt install -y libopenblas-base
      pip3 install torchvision
      pip3 uninstall -y torch
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

  4. Téléchargez le dépôt HuggingFace et les conditions requises pour l'installation.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone http://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone http://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. Téléchargez les configurations du modèle pré-entraîné.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
      gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

Entraîner le modèle

Entraînez le modèle à 2 milliards d'octets en utilisant une taille de lot de 16.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Résultat de l'analyse comparative

Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits des benchmarks pour différents types d'accélérateurs.

V5Litepod-4 V5Litepod-16 V5Litepod-64
Époque 3 3 3
config 600M 2 milliards 16 Mrds
Taille du lot global 64 128 256
Débit (exemples/s) 66 77 31

Entraîner ViT sur v5e

Ce tutoriel explique comment exécuter VIT dans la version v5e à l'aide du dépôt HuggingFace sur PyTorch/XLA sur l'ensemble de données cifar10.

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. Créer une ressource TPU:

    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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource QueuedResource sera à l'état ACTIVE:

     gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:

     state: ACTIVE

  3. Installer les dépendances Torch/xla

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

  4. Téléchargez le dépôt HuggingFace et les conditions requises pour l'installation.

       gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="
      git clone http://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget http://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"

Entraîner le modèle

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

Résultat de l'analyse comparative

Le tableau suivant présente les débits de benchmark pour les différents types d'accélérateurs.

V5Litepod-4 V5Litepod-16 V5Litepod-64
Époque 3 3 3
Taille du lot global 32 128 512
Débit (exemples/s) 201 657 2 844

TensorFlow 2.x

Entraîner Resnet sur un seul hôte v5e

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-4 ou v5litepod-8 à l'aide d'un ensemble de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez la page Préparer l'ensemble de données.

Préparation

  1. Créez des variables d'environnement :

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5litepod-4
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE peut être v5litepod-4 ou v5litepod-8. /

  2. Créer une ressource TPU:

    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} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE. Pour vérifier l'état de votre ressource en file d'attente, utilisez la commande suivante:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Se connecter à TPU à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Définir des variables d'environnement

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. Accédez au répertoire du dépôt de modèles et accédez aux exigences d'installation.

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

Entraîner le modèle

Exécutez le script d'entraînement.

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Supprimer le TPU et la ressource en file d'attente

  1. Supprimer votre TPU

    gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. Supprimer votre demande de ressource en file d'attente

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Entraîner Resnet sur un multi-hôte v5e

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-16 ou version ultérieure à l'aide d'un ensemble de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez Préparer l'ensemble de données.

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    La valeur de ACCELERATOR_TYPE peut être égale ou supérieure à v5litepod-16.

  2. Créer une ressource TPU:

    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} \
   --${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE. Pour vérifier l'état de votre ressource en file d'attente, utilisez la commande suivante:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Se connecter à votre TPU (nœud de calcul 0) à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Définir des variables d'environnement

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

  5. Accédez au répertoire du dépôt de modèles et accédez aux exigences d'installation.

     cd $MODELS_REPO && git checkout r2.15.0
 pip install -r official/requirements.txt

Entraîner le modèle

Exécutez le script d'entraînement.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

Supprimer le TPU et la ressource en file d'attente

  1. Supprimer votre TPU

    gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. Supprimer votre demande de ressource en file d'attente

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet