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

Entraînement Cloud TPU v5p

Cloud TPU v5p est le Cloud TPU de cinquième génération de Google Cloud, qui succède au TPU v4. v5p est optimisé pour l'entraînement à grande échelle et constitue une plate-forme de premier plan pour le développement de LLM, de modèles de diffusion et d'IA générative de base. De manière générale, v5p offre des performances jusqu'à deux fois supérieures à celles de la version 4, tout en empaquetant deux fois plus de TPU dans un pod (la plus grande tranche de 6 000 pixels contre 3 000 dans la version 4), ce qui se traduit par des performances jusqu'à quatre fois supérieures au niveau du pod. Il fonctionne également à une fréquence d'horloge plus élevée (1,75 GHz contre 1,05 GHz), ajoute SparseCore pour les représentations vectorielles continues à grande échelle et triple la capacité de la mémoire à haut débit (HBM).

Concepts Cloud TPU v5p

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

Les concepts de Cloud TPU (par exemple, les tranches, les hôtes et les TensorCores) et l'architecture système de Cloud TPU pour toutes les versions de 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 ou l'inférence. Ces types d'accélérateurs sont décrits dans les configurations v5p.

Gérer les ressources TPU

Toutes les commandes que vous pouvez utiliser pour gérer vos VM TPU sont décrites dans la section Gérer les TPU ou le Guide de l'utilisateur des ressources en file d'attente pour gérer les ressources en file d'attente.

Configuration du framework

Cette section décrit le processus de configuration général pour l'entraînement de modèle à l'aide de JAX ou PyTorch avec TPU v5p.

Configuration pour JAX

Si vos formes de secteurs comportent plus de quatre puces, plusieurs VM seront regroupées dans une tranche. Dans ce cas, vous devez utiliser l'option --worker=all pour exécuter l'installation sur toutes les VM TPU à l'aide d'une seule commande:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.20" -f https://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 v5p-32). 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.]

Utilisez --node=all pour exécuter la commande sur tous les nœuds de calcul multi-segments.

gcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Essayez les tutoriels JAX de ce document pour faire vos premiers pas avec l'entraînement v5p à l'aide de JAX.

Configuration pour PyTorch

L'environnement d'exécution PJRT est le seul compatible avec la version 5p, et 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 sur les pods v5p avec PyTorch/XLA 2.2.0 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
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Utilisez un script Python avec PJRT pour valider l'installation et afficher les appareils TPU disponibles (les sorties présentées ici ont été générées avec une tranche v5p-32).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

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...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Utilisez --node=all pour exécuter la commande sur tous les nœuds de calcul multi-segments.

gcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Suivez les tutoriels PyTorch de ce document pour faire vos premiers pas avec l'entraînement v5p à l'aide de PyTorch.

Surveiller et profiler

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

Tutoriels de formation

Cette section se concentre sur les tutoriels d'entraînement pour une seule tranche. Vous pouvez adapter ces tutoriels à l'entraînement à plusieurs tranches en ajoutant l'indicateur --node=all aux commandes SSH. Pour en savoir plus et consulter les bonnes pratiques, reportez-vous à l'introduction aux secteurs à secteurs multiples.

Diffusion 2.1
MaxText
ResNet sur un seul hôte v5p
ResNet sur une version v5p à hôtes multiples

Tutoriels JAX

Entraîner la diffusion 2.1

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

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

Créez des variables d'environnement :

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
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

Description des options de commande

Variable	Description
PROJECT_ID	Nom du projet Google Cloud
ACCELERATOR_TYPE	Consultez la page Versions de TPU pour connaître la version de votre TPU.
ZONE	Consultez le document Régions et zones de TPU pour connaître les zones compatibles.
RUNTIME_VERSION	Pour la version 5p, utilisez la version v2-alpha-tpuv5 pour RUNTIME_VERSION.
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: tpu-service-account@myprojectID.iam.gserviceaccount.com
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}
```
Remarque :L'option QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut est on-demand. Consultez la section quotas pour en savoir plus sur les différents types de quotas acceptés par Cloud TPU.

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE. Vérifiez l'état de votre ressource en file d'attente en exécutant la commande suivante:
```
gcloud alpha 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
```

Installez JAX et ses dépendances.

# compatible with v5p: only jax version 0.4.19 and later \
# jax 0.4.19 requires py 3.10 \

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

Téléchargez le dépôt HuggingFace et les exigences d'installation.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow 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='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'

Effectuer un nettoyage

À la fin de votre session, supprimez le TPU et la demande de ressources en file d'attente, ou pour supprimer les demandes de ressources en file d'attente dont l'état est "ÉCHEC". Pour supprimer une ressource en file d'attente, supprimez la ou les tranches, puis la demande de ressource en file d'attente en deux étapes:

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

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

Vous pouvez également utiliser --force pour supprimer les tranches et la requête de ressource en file d'attente en une seule étape:

# With --force
gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Résultats du benchmark

Le script d'entraînement "Diffusion stable" a été exécuté sur les versions v5p-8, v5p-32 et v5p-128. Le tableau suivant indique le débit.

	V5P-8	V5P-32	V5P-128
Étapes d'entraînement	150	150	150
Taille du lot global	32	64	64
Débit (exemples/s)	12.10	18,08	19,10

MaxText

Ce tutoriel explique comment entraîner le modèle MaxText à l'aide d'un ensemble de données synthétique sur Cloud TPU.

MaxText est un LLM hautes performances, arbitrairement évolutif, Open Source et bien testé, écrit en Python/JAX pur et ciblant les Cloud TPU. MaxText offre aux chercheurs et aux développeurs un outil accessible et adaptable pour repousser les limites de la recherche et du développement sur le traitement du langage naturel (TLN).

Avant d'exécuter ce tutoriel, vous devez configurer votre environnement Cloud TPU.

Configurer des variables d'environnement

export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name # user defined TPU name
export ACCELERATOR_TYPE=v5p-256
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export RUN_NAME=your_experiment_run_name # user defined name for this run
export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

Description des options de commande

Variable	Description
PROJECT_ID	Nom du projet Google Cloud
TPU_NAME	Nom défini par l'utilisateur pour votre TPU.
ACCELERATOR_TYPE	Consultez la page Versions de TPU pour connaître la version de votre TPU.
ZONE	Consultez le document Régions et zones de TPU pour connaître les zones compatibles.
RUNTIME_VERSION	Pour la version v5p, utilisez la version d'exécution "v2-alpha-tpuv5".
RUN_NAME	Nom d'exécution du test fourni par l'utilisateur.

Configuration facultative recommandée pour les multitranches:

export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

Si vous exécutez des charges de travail avec plusieurs tranches et que vous souhaitez des performances réseau optimales, envisagez de créer un réseau dédié avec une unité de transmission maximale (MTU) de 8 896 octets et de configurer des règles de pare-feu appropriées. Bien que facultative, cette étape peut améliorer considérablement les performances, en particulier lors du scaling à la hausse du nombre de tranches sur le réseau du centre de données (DCN). Notez que la création d'un réseau nécessite l'autorisation compute.networks.create dans le projet. Les exemples suivants montrent comment créer un réseau et une règle de pare-feu dédiés.

Créez un réseau dédié:

gcloud compute networks create ${NETWORK_NAME} \
--mtu=8896 \
--project=${PROJECT_ID} \
--subnet-mode=auto \
--bgp-routing-mode=regional

Créez une règle de pare-feu :

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}

Cloner le dépôt MaxText

git clone https://github.com/google/maxtext.git

Entraîner le modèle

Les sections suivantes décrivent deux options d'entraînement de MaxText.

Option 1

Si vous souhaitez qu'un script gère l'ensemble du workflow, du provisionnement des Cloud TPU à l'installation des dépendances, en passant par l'exécution de votre modèle et la suppression des ressources, vous pouvez utiliser multihost_job.py.

cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
--NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
--VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
--BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
--COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

Une fois le script lancé, un message semblable au suivant doit s'afficher dans le journal. L'emplacement du journal est référencé dans le message de sortie. Cliquez sur le premier lien pour accéder aux journaux de tous les nœuds de calcul une fois le provisionnement du TPU terminé.

------------------------------------

multihost_job finished running, TPUs are starting up to run your job remotely.

Logs for your job are displayed here:
https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID

To see the output of a single host, you may edit the slice and worker
number in the `log_file_path` property here:

https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID

When your job is finished, the main command log is in your Cloud Storage
bucket:
https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID

View the status of the created TPUs using:
gcloud alpha compute tpus queued-resources list --filter=RUN_NAME
--zone=ZONE --project=PROJECT_ID

Option 2

Pour exécuter le script d'entraînement plusieurs fois sur un Cloud TPU provisionné, utilisez le script multihost_runner.py pour utiliser la ressource.

Configurez des variables pour créer un TPU.
```
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export VALID_DURATION=1d
export QUOTA_TYPE=quota_type
```
Remarque :Les options suivantes ne sont nécessaires que si vous utilisez les multitranches:
```
--node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID
```
Remarque :Si vous spécifiez --node-count et --node-prefix, ne spécifiez pas également --node-id lorsque vous créez la ressource TPU à l'étape suivante. Une erreur sera générée si les trois indicateurs sont spécifiés.
Créez 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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}
```
Remarque :L'option QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut est on-demand. Consultez la section quotas pour en savoir plus sur les différents types de quotas acceptés par Cloud TPU.

Vous pourrez vous connecter à vos VM TPU à l'aide de SSH une fois que votre QueuedResource aura l'état ACTIVE:

Utilisez la commande describe pour interroger l'état de votre ressource en file d'attente.
```
gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE}
```
Lorsque votre ressource en file d'attente est à l'état ACTIF, le résultat ressemble à ce qui suit:
```
 state: ACTIVE
```

Se connecter à TPU à l'aide de SSH

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

Installer des dépendances

Remarque :Une fois que vous avez utilisé ssh pour vous connecter à votre VM TPU, vous devez définir les variables ${TPU_NAME} et ${MAXTEXT_OUTPUT_PATH} pour les utiliser dans la VM TPU:
```
export TPU_NAME=your_tpu_name
export MAXTEXT_OUTPUT_PATH=output-path
```
```
cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND='bash setup.sh'
```
Exécutez le modèle avec différents scripts de configuration, tels que 32b.sh et 64b.sh. Si vous exécutez le script à partir d'une VM TPU, vous devez ajouter l'option --INTERNAL_IP=true.
```
python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"
```
Remarque :Définissez --TPU_PREFIX=${NODE_PREFIX} pour "Multislice" au lieu de --TPU_PREFIX=${TPU_NAME}.

Effectuer un nettoyage

Supprimez votre TPU et vos ressources en file d'attente.

Résultats du benchmark

Le script d'entraînement MaxText a été exécuté de 32 à 1 160 milliards avec une précision bf16. Les résultats de ces exécutions sont présentés dans le tableau suivant.

Nombre de paramètres	Type d'accélérateur	TFLOP/puce/s	Utilisation des flops du modèle (MFU)
32 Mrds	V5P-128	3.28E+02	71,47%
64 Mrds	V5P-128	3.23E+02	70,31%
128 Mrds	V5P-256	3.15E+02	68,68%
128 Mrds	V5P-512	3.15E+02	68,53%
256 Mrds	V5P-1024	3.16E+02	68,82%
512 Mrds	V5P-1024	2,94E + 02	63,99%
1 024 milliards	V5P-2048	2,49E + 02	64,05%
1 024 milliards	V5P-4096	2,97E + 02	64,80%
1 160 milliards	V5P-7680	2,95E + 02	64,27%
1 160 milliards	V5P-12288	3,04E + 02	66,23%

Le modèle de paramètre 256 milliards a été testé sur les versions v5p-512 et v5p-1024 en utilisant à la fois les pondérations bf16 et int8. Le tableau suivant affiche les résultats de ces tests.

	V5P-512	V5P-512	V5P-1024	V5P-1024
Taille du lot global (jetons)	5,24 E + 05	5,24 E + 05	1,05E + 06	1,05E + 06
Précision	bf16	int8	bf16	int8
TFLOP/puce/s	307	408	308	414
Utilisation des flops du modèle (MFU)	66,98%	88,85%	67,09%	90,23%

Tutoriels TensorFlow

Entraîner ResNet sur un seul hôte v5p

Ce tutoriel explique comment entraîner ImageNet sur un TPU v5p-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

Créez des variables d'environnement :

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

Pour ce tutoriel, utilisez v5p-8 comme ACCELERATOR_TYPE.

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}
```
Remarque :L'option QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut est on-demand. Consultez la section quotas pour en savoir plus sur les différents types de quotas acceptés par Cloud TPU.

Vous pourrez vous connecter à votre VM TPU à l'aide de SSH 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 alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}
```

Se connecter à TPU à l'aide de SSH

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

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

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"

Effectuer un nettoyage

Supprimez votre TPU et vos ressources en file d'attente.

Entraîner ResNet sur une version v5p à hôtes multiples

Ce tutoriel explique comment entraîner ImageNet sur v5p-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.

Créez des variables d'environnement :

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

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

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}
```
Remarque :L'option QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut est on-demand. Consultez la section quotas pour en savoir plus sur les différents types de quotas acceptés par Cloud TPU.

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

Utilisez la commande describe pour interroger l'état de votre ressource en file d'attente:
```
gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}
```

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}

Définir des variables d'environnement

export TPU_NAME=your_tpu_name
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

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 \
  --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"

Effectuer un nettoyage

Supprimez votre TPU et vos ressources en file d'attente.

PyTorch/XLA

Llama 2

Ce tutoriel explique comment entraîner le modèle Llama 2 7B en v5p à l'aide d'une copie du dépôt HuggingFace sur PyTorch/XLA avec un chargement en parallèle général et évolutif pour les graphiques de calcul ML (GSPMD).

Préparation

Créez des variables pour l'ID de projet, le type d'accélérateur, la zone, la version d'exécution et le nom du TPU.

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
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_DURATION=1d

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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}
```
Remarque :L'option QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut est on-demand. Consultez la section quotas pour en savoir plus sur les différents types de quotas acceptés par Cloud TPU.

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

Utilisez la commande describe pour interroger l'état de votre ressource en file d'attente.
```
gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE}
```
Lorsque votre ressource en file d'attente est à l'état ACTIF, le résultat ressemble à ce qui suit:
```
 state: ACTIVE
```
Remarque :Si vous exécutez un environnement à plusieurs tranches, utilisez gcloud alpha compute tpus queued-resources ssh --node=all ... au lieu de gcloud compute tpus tpu-vm ssh ... pour vos commandes. Cela garantit que toutes vos puces TPU sont correctement utilisées.

Installez Pytorch/XLA et les dépendances requises.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip3 install typing-extensions
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

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 -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install git+file://$PWD
pip3 install datasets accelerate evaluate scikit-learn'

Téléchargez la configuration du modèle 7B.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

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
export XLA_USE_BF16=1
export XLA_IR_DEBUG=1
export XLA_HLO_DEBUG=1

export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_gather=true
--xla_jf_spmd_threshold_for_windowed_einsum_mib=0"

export PROFILE_EPOCH=0
export PROFILE_STEP=3
export PROFILE_DURATION_MS=20000
export PROFILE_LOGDIR=/tmp/home/

cd transformers
python examples/pytorch/language-modeling/run_clm.py \
 --tokenizer_name hf-internal-testing/llama-tokenizer \
 --dataset_name wikitext \
 --dataset_config_name wikitext-2-raw-v1 \
 --per_device_train_batch_size 96 \
 --per_device_eval_batch_size 8 \
 --num_train_epochs 1 \
 --do_train \
 --output_dir /tmp/output \
 --overwrite_output_dir \
 --config_name ~/config.json \
 --save_strategy no \
 --logging_strategy no \
 --remove_unused_columns no \
 --optim adafactor \
 --torch_dtype bfloat16 \
 --dataloader_drop_last yes \
 --block_size 2048 \
 --spmd_2d_sharding 1 \
 --spmd_grad_chkpt
'

Si vous exécutez un environnement à plusieurs tranches, vous devez définir l'option --spmd_dcn_parallelism sur le nombre de tranches.

Le SPMD_USER_GUIDE (guide de l'utilisateur) fournit un guide utilisateur plus détaillé qui explique les différentes variables d'environnement et les boutons d'activation/de désactivation du script HF. Notez que LIBTPU_INIT_ARGS sera intégré à PyTorch/XLA et activé par défaut dans les prochaines versions.

Effectuer un nettoyage

Supprimez votre TPU et vos ressources en file d'attente.

Résultats du benchmark

Le débit des trois tailles de modèle Llama 2 est inclus dans le tableau suivant.

	V5P-8	V5P-128	V5P-128
Taille du modèle	7 Mrds	13 Mrds	70 Mrds
Taille du lot global	96	1 024	128
Forme du maillage de segmentation	(4, 1)	(64, 1).	(16, 4).
Utilisation des flops du modèle (MFU)	56,67%	55,80%	51,85%

Assistance et commentaires

Tous vos commentaires sont les bienvenus ! Pour partager vos commentaires ou demander de l'aide, remplissez le formulaire d'assistance ou de commentaires pour Cloud TPU.